https://zhuanlan.zhihu.com/p/191569603

[toc]

C++部署pytorch模型

前言

项目需要将pytorch训练好的网络用c++调用，在正式开始项目之前，在网上查了各种资料，共有三种实现方法： 直接将网络从最基础的CNN模块用C++实现； 将网咯模型和参数保存，然后使用opencv的DNN模块加载，这个方法tensorflow、torch等其他网络架构也能用，具体包含哪些下文会给出； * 使用pytorch官网提供的c++接口：LibTorch。其原理也是保存网络模型和参数，然后用LibTorch进行加载。 由于第一项c++从第层撸起太过硬核，自己水平有限，此处就不做介绍。大佬们可以自己尝试。此处只介绍opencv和LibTorch实现的方法。

运行环境：
win10 64位
cuda 10.2
pytorch 1.6.0
torchvision 0.7
opencv 4.3
vs2019
LibTorch 1.6
ps:
pytorch相关软件都是直接在官网下载的最新版本。

训练一个简单的pytorch网络

首先，参考pytorch官方文档中训练一个分类器的代码，训练一个简单的图像分类器。代码如下：

import torch.optim as optim
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.onnx
import torchvision
import torchvision.transforms as transforms
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root=’./data’, train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(root=’./data’, train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
shuffle=False, num_workers=0)

classes = (‘plane’, ‘car’, ‘bird’, ‘cat’,
‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’)

# functions to show an image

def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()

# get some random training images
dataiter = iter(trainloader)
images, labels = dataiter.next()

print(images.shape)

# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(’ ‘.join(’%5s’ % classes[labels[j]] for j in range(4)))

class Net(nn.Module):
def init(self):
super(Net, self).init()
self.conv1 = nn.Conv2d(3, 6, 3)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 12, 3)
self.conv3 = nn.Conv2d(12, 32, 3)
self.fc1 = nn.Linear(32 4 4, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)

<span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
    <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">conv1</span><span class="p">(</span><span class="n">x</span><span class="p">)))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">conv2</span><span class="p">(</span><span class="n">x</span><span class="p">)))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">conv3</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">32</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">*</span> <span class="mi">4</span><span class="p">)</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">fc1</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">fc2</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fc3</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">x</span>

net = Net()
net.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(100): # loop over the dataset multiple times

<span class="n">running_loss</span> <span class="o">=</span> <span class="mf">0.0</span>
<span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">data</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">trainloader</span><span class="p">,</span> <span class="mi">0</span><span class="p">):</span>
    <span class="c1"># get the inputs; data is a list of [inputs, labels]</span>
    <span class="n">inputs</span><span class="p">,</span> <span class="n">labels</span> <span class="o">=</span> <span class="n">data</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">),</span> <span class="n">data</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="c1"># zero the parameter gradients</span>
    <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>

    <span class="c1"># forward + backward + optimize</span>
    <span class="n">outputs</span> <span class="o">=</span> <span class="n">net</span><span class="p">(</span><span class="n">inputs</span><span class="p">)</span>

    <span class="n">loss</span> <span class="o">=</span> <span class="n">criterion</span><span class="p">(</span><span class="n">outputs</span><span class="p">,</span> <span class="n">labels</span><span class="p">)</span>
    <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
    <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>

    <span class="c1"># print statistics</span>
    <span class="n">running_loss</span> <span class="o">+=</span> <span class="n">loss</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">i</span> <span class="o">%</span> <span class="mi">2000</span> <span class="o">==</span> <span class="mi">1999</span><span class="p">:</span>    <span class="c1"># print every 2000 mini-batches</span>
        <span class="k">print</span><span class="p">(</span><span class="n">outputs</span><span class="p">)</span>
        <span class="k">print</span><span class="p">(</span><span class="s1">'[</span><span class="si">%d</span><span class="s1">, </span><span class="si">%5d</span><span class="s1">] loss: </span><span class="si">%.3f</span><span class="s1">'</span> <span class="o">%</span>
              <span class="p">(</span><span class="n">epoch</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="n">running_loss</span> <span class="o">/</span> <span class="mi">2000</span><span class="p">))</span>
        <span class="n">running_loss</span> <span class="o">=</span> <span class="mf">0.0</span>

print(‘Finished Training’)

上述代码相对于官方文档的代码，仅仅是增加了卷积层和利用GPU进行训练，且输出结果未经处理，只是简单输出各个类别的概率值。 训练完网络之后，将网络保存，代码如下：

# 保存网络结构和参数

# 方法1：保存网络结构和参数
PATH = ‘./cifar_net.pth’
torch.save(net, PATH)

# 方法2：保存网络参数
PATH = ‘./cifar_net.pth’
torch.save(net.state_dict(), PATH)

# 方法3：导出网络到ONNX
dummy_input = torch.randn(1, 3, 32, 32).to(device)
torch.onnx.export(net, dummy_input, “torch.onnx”)

# 方法4：保存网络位TORCHSCRIPT
dummy_input = torch.randn(1, 3, 32, 32).to(device)
traced_cell = torch.jit.trace(net, dummy_input)
traced_cell.save(“tests.pth”)

上述四种保存方法本文主要使用方法3和方法4，具体应用方式在下文会详细说明。此处简单说一下一个比较坑的地方：我开始以为使用方法1保存的网络可以像tensorflow那样直接用load函数导入，自动重建出原始网络架构，但是试验后才发现，要想成功导入，还需要将定义网络的类也放在对应的py文件中，这有点。。。 方法1导入示例代码如下：

import torch

import torchvision

import torch.optim as optim

import torch.nn.functional as F

import torch.nn as nn

import cv2

import torchvision.transforms as transforms

PATH = ‘./cifar_net.pth’

device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)

class Net(nn.Module):

def init(self):

super(Net, self).init()

self.conv1 = nn.Conv2d(3, 6, 3)

self.pool = nn.MaxPool2d(2, 2)

self.conv2 = nn.Conv2d(6, 12, 3)

self.conv3 = nn.Conv2d(12, 32, 3)

self.fc1 = nn.Linear(32  4  4, 120)

self.fc2 = nn.Linear(120, 84)

self.fc3 = nn.Linear(84, 10)
<span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
    <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">conv1</span><span class="p">(</span><span class="n">x</span><span class="p">)))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">conv2</span><span class="p">(</span><span class="n">x</span><span class="p">)))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">conv3</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">32</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">*</span> <span class="mi">4</span><span class="p">)</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">fc1</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">fc2</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fc3</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">x</span>

model = torch.load(PATH)

到此，所有的准备工作已经完成。

利用opencv加载训练好的模型和网络

参考链接： OpenCV4.0 运行快速风格迁移（Torch） opencv官方文档 OpenCV加载Pytorch模型出现Unsupported Lua type 解决方法 将模型从PYTORCH导出到ONNX并使用ONNX RUNTIME运行（官网链接） 这个方法算是一个比较常用的方法，而且可以用到许多深度学习框架上面： 根据opencv官方文档中的说明，可以支持以下框架：Caffe，Darknet，Onnx，Tensorflow，Torch等。但是很可惜，没有我用的pytoch，但是根据第三个参考链接中的方法，可以利用ONNX实现曲线救国。首先利用保存模型方法3所示的办法，将网络和参数保存为对应的格式。然后使用opencv提供的Net cv::dnn::readNetFromONNX ( const String & onnxFile )函数读取保存好的网络。代码实现如下：

//测试opencv加载pytorch模型

#include <opencv2/dnn.hpp>

#include <opencv2/imgproc.hpp>

#include <opencv2/highgui.hpp>

using namespace cv;

using namespace cv::dnn;

#include <fstream>

#include <iostream>

#include <cstdlib>

using namespace std;

int main()
{
String modelFile = “./torch.onnx”;
String imageFile = “./dog.jpg”;

<span class="n">dnn</span><span class="o">::</span><span class="n">Net</span> <span class="n">net</span> <span class="o">=</span> <span class="n">cv</span><span class="o">::</span><span class="n">dnn</span><span class="o">::</span><span class="n">readNetFromONNX</span><span class="p">(</span><span class="n">modelFile</span><span class="p">);</span> <span class="c1">//读取网络和参数

Mat image = imread(imageFile); // 读取测试图片
cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
Mat inputBolb = blobFromImage(image, 0.00390625f, Size(32, 32), Scalar(), false, false); //将图像转化为正确输入格式

net.setInput(inputBolb); //输入图像

Mat result = net.forward(); //前向计算

cout << result << endl;
}

上述代码就是对第一个参考链接的代码进行了简化，且将输入网络的模型从torch改成ONNX格式。 运行结果如下：

[-0.19793352, -4.0697966, 1.2769811, 2.7011304, 0.22390884, 1.9039617, -0.47333384, -0.15912014, 0.32441139, -2.4327304]

如果需要部署其他深度学习框架的网络，执行步骤基本类似。

利用pytorch官方提供的LibTorch加载训练好的模型和网络

参考链接： windows+VS2019+PyTorchLib配置使用攻略 C++调用pytorch，LibTorch在win10下的vs配置和cmake的配置 在C ++中加载TORCHSCRIPT模型官网链接 此处首先说明一下将pytroch保存为TORCHSCRIPT的方法有两种，一种是追踪式，另一种是脚本式。具体介绍见官方文档，理论上此方法两种保存方式都行，方法4中的是追踪式的方法，此文使用此方法。 首先按照第一个参考链接中的方法配置LibTorch环境，然后复制粘贴其中的示例代码，进行测试，但是我个人在运行的时候ToTensor(image).to(at::kCUDA);这个语句报错了，提示未定义ToTensor()，这句话的功能也很简单，就是将普通图像格式转化为模型输入需要的格式，于是我又根据第二个参考链接将转化代码进行了修改，代码如下：

#include <torch/script.h>

#include <iostream>

#include <opencv2/opencv.hpp>

#include <torch/torch.h>



// 有人说调用的顺序有关系，我这好像没啥用~~



int main()

{

torch::DeviceType device_type;

if (torch::cuda::is_available()) {

std::cout << “CUDA available! Predicting on GPU.” << std::endl;

device_type = torch::kCUDA;

}

else {

std::cout << “Predicting on CPU.” << std::endl;

device_type = torch::kCPU;

}

torch::Device device(device_type);
<span class="c1">//Init model

std::string model_pb = “tests.pth”;
auto module = torch::jit::load(model_pb);
module.to(at::kCUDA);

<span class="k">auto</span> <span class="n">image</span> <span class="o">=</span> <span class="n">cv</span><span class="o">::</span><span class="n">imread</span><span class="p">(</span><span class="s">"dog.jpg"</span><span class="p">,</span> <span class="n">cv</span><span class="o">::</span><span class="n">ImreadModes</span><span class="o">::</span><span class="n">IMREAD_COLOR</span><span class="p">);</span>
<span class="n">cv</span><span class="o">::</span><span class="n">Mat</span> <span class="n">image_transfomed</span><span class="p">;</span>
<span class="n">cv</span><span class="o">::</span><span class="n">resize</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">image_transfomed</span><span class="p">,</span> <span class="n">cv</span><span class="o">::</span><span class="n">Size</span><span class="p">(</span><span class="mi">32</span><span class="p">,</span> <span class="mi">32</span><span class="p">));</span>

<span class="c1">// convert to tensort

torch::Tensor tensor_image = torch::from_blob(image_transfomed.data,
{ image_transfomed.rows, image_transfomed.cols,3 }, torch::kByte);
tensor_image = tensor_image.permute({ 2,0,1 });
tensor_image = tensor_image.toType(torch::kFloat);
tensor_image = tensor_image.div(255);
tensor_image = tensor_image.unsqueeze(0);
tensor_image = tensor_image.to(at::kCUDA);
torch::Tensor output = module.forward({ tensor_image }).toTensor();
auto max_result = output.max(1, true);
auto max_index = std::get<1>(max_result).item<float>();
std::cout << output << std::endl;
//return max_index;
return 0;

}

运行结果如下：

CUDA available! Predicting on GPU.

1.0824 -4.6106  1.0189  2.9937  1.4570  1.4964 -1.3164 -0.7753  0.4567 -3.2543

[ CUDAFloatType{1,10} ]