Pytorch model deployment ------- pytorch deployment full record

Full record of Pytorch model deployment

pytorch model save

Model save

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 loading

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’)) # 仅加载参数

Save and load Checkpoint for inference/continue training

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

load

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()

Model conversion

libtorch does not depend on python, the model trained by python needs to be converted to script model before it can be loaded by 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)

Own model

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")

It should be noted that if you want to train the model on the gpu and do inference on the cpu , you must transform before the model save, and then remember to call model.eval(), which looks like

• That is to say, what the cpu saves can only be inference on the cpu, that is to say, what the gpu saves can only be inference on the gpu,

onnx model check

import onnx

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

Install libtorch

Download link

Write program to load model

#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}) 

Test (cpu and gpu)

You can choose to put data on cpu or 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} ]

Use Tensorrt to accelerate

Tensort accelerated link