Pytorch模型部署之-------pytorch部署全记录

Pytorch模型部署全纪录

pytorch模型保存

模型保存

torch.save(model, ‘net.pth’) # 保存整个神经网络的模型结构以及参数
torch.save(model, ‘net.pkl’) # 保存整个神经网络的模型结构以及参数
torch.save(model.state_dict(), ‘net_params.pth’) # 只保存模型参数
torch.save(model.state_dict(), ‘net_params.pkl’) # 只保存模型参数

模型加载

model = torch.load(‘net.pth’) # 加载整个神经网络的模型结构以及参数
model = torch.load(‘net.pkl’) # 加载整个神经网络的模型结构以及参数
model.load_state_dict(torch.load(‘net_params.pth’)) # 仅加载参数
model.load_state_dict(torch.load(‘net_params.pkl’)) # 仅加载参数

保存和加载 Checkpoint 用于推理/继续训练

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

加载

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

模型转换

libtorch不依赖于python，python训练的模型，需要转换为script model才能由libtorch加载

import torch
import architecture as arch
 
# An instance of your model.
model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \
                        mode='CNA', res_scale=1, upsample_mode='upconv')
 
model.load_state_dict(torch.load('./models/RRDB_ESRGAN_x4.pth'), strict=True)
model.eval()
 
# An example input you would normally provide to your model's forward() method.
example = torch.rand(64, 3, 3, 3)
 
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
output = traced_script_module(torch.ones(64, 3, 3, 3))
traced_script_module.save("./models/RRDB_ESRGAN_x4_000.pt")
 
# The traced ScriptModule can now be evaluated identically to a regular PyTorch module
print(output)

自己的模型

import torch
import torchvision
from PIL import Image
import numpy as np
from model import AlexNet

# 图片发在了build文件夹下
image = Image.open("/home/zhongsy/datasets/dataset/train/no_obstacle/0.jpg")
image = image.resize((224, 224), Image.ANTIALIAS)
image = np.asarray(image)
image = image / 255
image = torch.Tensor(image).unsqueeze_(dim=0)
image = image.permute((0, 3, 1, 2)).float()

model = torch.load('./AlexNet.pt',map_location=torch.device('cpu'))
model.eval()
input_cpu_ = image.cpu()
input_gpu = image.cuda()


torchd_cpu = torch.jit.trace(model, input_cpu_)
torch.jit.save(torchd_cpu, "cpu.pth")

model_gpu = torch.load('./AlexNet.pt')
model_gpu.cuda()
model_gpu.eval()

torchd_gpu = torch.jit.trace(model_gpu,input_gpu)
torch.jit.save(torchd_gpu, "gpu.pth")

需要注意的是如果想要在gpu上训练模型，在cpu上做inference，一定要在模型save之前转化，再就是记得调用model.eval(),形如

• 也就是说 cpu保存的只能在cpu上做推理,也就是说 gpu保存的只能在gpu上做推理,

onnx模型检查

import onnx

model = onnx.load('gpu.onnx')
onnx.checker.check_model(model)
print("====> pass")

安装libtorch

下载链接

编写程序加载模型

#include <time.h>
#include <torch/script.h>

#include <algorithm>
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>

int main() {
  clock_t start, end;
  torch::jit::script::Module model = torch::jit::load("../cpu.pth");
//   model.to(at::kCUDA);
  start = clock();

  cv::Mat input_image = cv::imread("../294.jpg");
  cv::resize(input_image, input_image, cv::Size(224, 224));

  torch::Tensor image_tensor = torch::from_blob(
      input_image.data, {input_image.rows, input_image.cols, 3}, torch::kByte);

  image_tensor = image_tensor.permute({2, 0, 1});

  image_tensor = image_tensor.toType(torch::kFloat);

  image_tensor = image_tensor.div(255);

  image_tensor = image_tensor.unsqueeze(0);

//   image_tensor = image_tensor.to(at::kCUDA);

  torch::Tensor pred = model.forward({image_tensor}).toTensor();
  end = clock();
  std::cout << "F2运行时间" << (double)(end - start) / CLOCKS_PER_SEC
            << std::endl;

  std::cout << pred << std::endl;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.4 FATAL_ERROR)
project(simnet)

set(Torch_DIR /home/zhongsy/Downloads/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)        # 查找libtorch
find_package(OpenCV REQUIRED)       # 查找OpenCV
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
if(NOT Torch_FOUND)
    message(FATAL_ERROR "Pytorch Not Found!")
endif(NOT Torch_FOUND)

message(STATUS "Pytorch status:")
message(STATUS "    libraries: ${TORCH_LIBRARIES}")

message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")

add_executable(simnet pytorch.cc)
target_link_libraries(simnet ${TORCH_LIBRARIES} ${OpenCV_LIBS}) 

测试(cpu与gpu)

可以把数据选择放在cpu或者gpu上面

cpu运行时间0.094049 s
 0.4806 -0.5249
[ CPUFloatType{1,2} ]

cuda::is_available():1
Time used:33.31 ms
-10.0399  10.7939
[ CUDAFloatType{1,2} ]

利用Tensorrt加速

tensort加速链接