【pytorch】将模型部署至生产环境：借助TorchScript跟踪法及注释法生成可供C++调用的模块

颢师傅

已于 2022-05-06 02:02:07 修改

阅读量1.8k

点赞数

分类专栏：计算机视觉 python pytorch 文章标签： pytorch python

于 2022-05-04 21:25:08 首次发布

本文链接：https://blog.csdn.net/hh1357102/article/details/124573299

版权

计算机视觉同时被 3 个专栏收录

132 篇文章 10 订阅

订阅专栏

python

89 篇文章 11 订阅

订阅专栏

pytorch

65 篇文章 3 订阅

订阅专栏

(一）思路简介
1.pyTorch会提供TorchScript，它可以生成最有效的C++可读的模型格式。
2.TorchScript的记录法支持Python控制流的子集，由TorchScript创建的中间阶段可以由LibTorch读取。
3.接下来，由LibTorch加载的模型可以保存为C++对象，并且可以在C++程序中运行。
对于TorchScript的注释法，可以从torch.jit.ScriptModule继承。这有助于PyTorch避免使用纯Python方法，而这些方法无法转换为TorchScript。
在这里插入图片描述
（二）待训练模型采用
CIFAR10，10分类
按上述源码训练后得到模型参数文件：saveTextOnlyParams.pth
（三）使用跟踪法将模型序列化：

import os.path
from typing import Iterator
import numpy as np
import torch
import cv2
from PIL import Image
from torch.utils.data import Dataset,DataLoader,Subset,random_split
import re
from functools import reduce
from torch.utils.tensorboard import SummaryWriter as Writer
from torchvision import transforms,datasets
import torchvision as tv
from torch import nn
import torch.nn.functional as F
import time
import onnx
import onnxruntime
#查看命令：tensorboard --logdir=./myBorderText
#可用pycharm中code中的generater功能实现：

class myCustomerNetWork(nn.Module):
    def __init__(self):
        super().__init__()
        #输入3通道输出6通道：
        self.features=nn.Sequential(nn.Conv2d(3, 64, (3, 3)),nn.ReLU(),nn.Conv2d(64,128,(3,3)),
                                    nn.ReLU(),nn.Conv2d(128,256,(3,3)),nn.ReLU(),nn.AdaptiveAvgPool2d(1))

        self.classfired=nn.Sequential(nn.Flatten(),nn.Linear(256,80),nn.Dropout(),nn.Linear(80,10))

    def forward(self,x):
        return self.classfired(self.features(x))
#网络输入要求为torch.Size([32, 3, 32, 32])格式

myNet=myCustomerNetWork()
pthfile = r'D:\flask_pytorch\saveTextOnlyParams.pth'
#当strict=false时，参数文件匹配得上就加载，没有就默认初始化。
myNet.load_state_dict(torch.load(pthfile),strict=False)
if torch.cuda.is_available():
    myNet=myNet.cuda()
myNet.eval()
if __name__ == '__main__':
    #对于opencv中的dnn模块，网络输入必须为(n,c,w,h)格式
    imagePath = r"C:\Users\25360\Desktop\monodepth.jpeg"
    img = cv2.imdecode(np.fromfile(imagePath, np.uint8), -1)
    img = cv2.resize(img, (32, 32))
    # bgr转rgb
    img = img[:, :, ::-1].copy()
    inputX = torch.FloatTensor(img).cuda()
    inputX = inputX.permute(2, 0, 1).contiguous()
    inputX = inputX.unsqueeze(0)
    with torch.no_grad():
        jit_model = torch.jit.trace(myNet, inputX)
    # 将模型序列化
    jit_model.save('jit_model.pth')

可用pytorch直接加载序列化模型并在python环境中使用：

if __name__ == '__main__':
    #网络输入为(n,c,w,h)格式
    imagePath = r"C:\Users\25360\Desktop\monodepth.jpeg"
    img = cv2.imdecode(np.fromfile(imagePath, np.uint8), -1)
    img = cv2.resize(img, (32, 32))
    # bgr转rgb
    img = img[:, :, ::-1].copy()
    inputX = torch.FloatTensor(img).cuda()
    inputX = inputX.permute(2, 0, 1).contiguous()
    inputX = inputX.unsqueeze(0)
    with torch.no_grad():
        jit_model = torch.jit.trace(myNet, inputX)
    # 将模型序列化
    #jit_model.save('jit_model.pth')
    jit_model = torch.jit.load('jit_model.pth')
    print(jit_model(inputX))

输出：

tensor([[  71.6224,   10.6505,  165.0648,  313.5768, -148.1144,  329.7959,
          109.9136, -266.1085, -171.0974, -272.6216]], device='cuda:0',
       grad_fn=<AddmmBackward0>)

（四）使用注释法将模型序列化：

import os.path
from typing import Iterator
import numpy as np
import torch
import cv2
from PIL import Image
from torch.utils.data import Dataset,DataLoader,Subset,random_split
import re
from functools import reduce
from torch.utils.tensorboard import SummaryWriter as Writer
from torchvision import transforms,datasets
import torchvision as tv
from torch import nn
import torch.nn.functional as F
import time
import onnx
import onnxruntime
#查看命令：tensorboard --logdir=./myBorderText
#可用pycharm中code中的generater功能实现：
#向模型添加显式注释：
class myCustomerNetWork(torch.jit.ScriptModule):
    def __init__(self):
        super().__init__()
        #输入3通道输出6通道：
        self.features=nn.Sequential(nn.Conv2d(3, 64, (3, 3)),nn.ReLU(),nn.Conv2d(64,128,(3,3)),
                                    nn.ReLU(),nn.Conv2d(128,256,(3,3)),nn.ReLU(),nn.AdaptiveAvgPool2d(1))

        self.classfired=nn.Sequential(nn.Flatten(),nn.Linear(256,80),nn.Dropout(),nn.Linear(80,10))

    @torch.jit.script_method
    def forward(self,x):
        return self.classfired(self.features(x))
#网络输入要求为torch.Size([32, 3, 32, 32])格式

myNet=myCustomerNetWork()
pthfile = r'D:\flask_pytorch\saveTextOnlyParams.pth'
#当strict=false时，参数文件匹配得上就加载，没有就默认初始化。
myNet.load_state_dict(torch.load(pthfile),strict=False)
if torch.cuda.is_available():
    myNet=myNet.cuda()
myNet.eval()


if __name__ == '__main__':
    # 将模型序列化
    myNet.save('jit_model2.pth')

可用pytorch直接加载序列化模型并在python环境中使用：

if __name__ == '__main__':
    #网络输入为(n,c,w,h)格式
    imagePath = r"C:\Users\25360\Desktop\monodepth.jpeg"
    img = cv2.imdecode(np.fromfile(imagePath, np.uint8), -1)
    img = cv2.resize(img, (32, 32))
    # bgr转rgb
    img = img[:, :, ::-1].copy()
    inputX = torch.FloatTensor(img).cuda()
    inputX = inputX.permute(2, 0, 1).contiguous()
    inputX = inputX.unsqueeze(0)
    # with torch.no_grad():
    #     jit_model = torch.jit.trace(myNet, inputX)
    # 将模型序列化
    #myNet.save('jit_model2.pth')
    jit_model = torch.jit.load('jit_model2.pth')
    print(jit_model(inputX))

输出：

tensor([[  71.6224,   10.6505,  165.0648,  313.5768, -148.1144,  329.7959,
          109.9136, -266.1085, -171.0974, -272.6216]], device='cuda:0',
       grad_fn=<AddmmBackward0>)

使用libtorch运行生成的pth文件，有3点注意事项：
1.要保证libtorch版本与pytorch版本保持一致
2.若出现c10::NotImplementedError,则在VS的链接中加入如下命令行：
/INCLUDE:“?ignore_this_library_placeholder@@YAHXZ”
如同所示：
在这里插入图片描述
libtorch所链接的lib有：
torch.lib
torch_cuda.lib
torch_cuda_cu.lib
torch_cuda_cpp.lib
torch_cpu.lib
c10_cuda.lib
caffe2_nvrtc.lib
c10.lib
kineto.lib
dnnl.lib
fbgemm.lib
asmjit.lib
XNNPACK.lib
cpuinfo.lib
clog.lib
libprotoc.lib
pthreadpool.lib
libprotobuf.lib
libprotobuf-lite.lib

C++调用代码为：

#include <iostream>
#include "torch/script.h"
#include<torch/torch.h>
#include <iostream>
#include<opencv.hpp>
int main()
{

    cv::Mat image = cv::imread("C:/Users/25360/Desktop/monodepth.jpeg", cv::IMREAD_COLOR);
    if (image.empty())
    {
        return 0;
    }
    std::string modelname = "jit_model2.pth";
    torch::jit::Module model = torch::jit::load(modelname);
    model.eval();
    model.to(at::kCUDA);
    //cv::namedWindow("image", cv::WindowFlags::WINDOW_NORMAL);
    //cv::imshow("image", image);
    //cv::waitKey(0);
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
    cv::resize(image, image, cv::Size(32, 32));
    torch::Tensor myTensor=torch::from_blob(image.data, { 1,32,32,3 }, torch::kByte);
    myTensor = myTensor.permute({ 0,3,1,2 });
    myTensor = myTensor.to(torch::kF32);
    myTensor = myTensor.to(torch::kCUDA);
    auto output = model.forward({ myTensor }).toTensor();
    std::cout << output;

return 0;

}

输出为：
Columns 1 to 8 71.6224 10.6505 165.0648 313.5768 -148.1144 329.7959 109.9136 -266.1085

Columns 9 to 10-171.0975 -272.6216
在这里插入图片描述
输出结果与pytorch直接调用结果一致。

颢师傅

关注

0
点赞
踩
2

收藏

觉得还不错? 一键收藏
打赏
0
评论
【pytorch】将模型部署至生产环境：借助TorchScript跟踪法及注释法生成可供C++调用的模块

pyTorch会提供TorchScript，它可以生成最有效的C++可读的模型格式。TorchScript可以将用Python编写的模型转换为高度优化的中间表示，然后LibTorch可以读取这些模型。由LibTorch加载的模型可由C++调用。1.要保证libtorch版本与pytorch版本保持一致2.若出现c10::NotImplementedError,则在VS的链接中加入如下命令行：/INCLUDE:"?ignore_this_library_placeholder@@YAHXZ"
复制链接

扫一扫