preface
上次在服务器上装上了 docker 版的 OpenCV,就是为本篇文章服务的,因为最近做模型的部署,就要用 C++ 调用 pytorch 的模型,pytorch 推出了 libtorch 来进行模型的部署,本篇文章就记录一下部署的过程。
转化模型
用 pytorch 训练完的模型必须要在 python 环境下才能够被调用,既然我们要让它能够被 C++ 调用那就得转化模型,pytorch 官方提出的 TorchScript 就是用来做这事的,TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。任何 TorchScript 程序都可以从 Python 进程中保存并在没有 Python 依赖项的进程中加载。
那么具体怎么做呢,也很简单,就是用 torch.jit.trace
对模型进行追踪,拿我们训练好的 mnist 模型来做例子(部分代码未给出,只有核心代码),torch.jit.trace
函数接收模型和示例输入这两个参数,其中示例输入的尺寸要与模型的输入尺寸相同,为 (1, channels, img_width, img_height)
,这里 LeNet 模型,输入的图像是单通道 28*28 大小的。这样模型就被追踪了,将其序列化为 .pt
后缀的模型文件,这样我们就离开了 python 的领域,准备跨入 C++ 领域了。
import torch
model = Net().to('cpu')
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
load_checkpoint('mnist-4690.pth', model, optimizer)
example = torch.rand(1, 1, 28, 28)
traced_model = torch.jit.trace(model, example)
output = traced_model(torch.ones(1, 1, 28, 28))
print(output)
traced_model.save('traced_mnist_pytorch_model.pt')
代码的输出为:tensor([[ -1.9272, -10.6155, -0.5130, 4.6338, -7.4411, 0.6665, -8.3521, -3.3505, 16.8189, 3.7142]], grad_fn=
)
下载依赖库
C++ 读图像大多用 OpenCV,我们已经编译好了,物理机编译可以参考我这篇教程,然后还要用到官方的 libtorch
库,直接去官网下载对应的版本然后解压就行了,不需要安装。
编写C++代码
接下去就利用 OpenCV 和 libtorch 对模型进行调用,注意了,测试图片的时候要将图片进行预处理,预处理步骤和在 PyTorch 中的预处理要一模一样,否则会出错。
#include <torch/script.h>
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>
using namespace std;
using namespace cv;
int main(int argc, char * argv[]){
if (argc != 2) {
std::cerr << "no module found !\n";
return -1;
}
torch::jit::script::Module module;
// 加载序列化的PyTorch模型
try {
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e){
std::cerr << "error loading the module\n";
return -1;
}
std::cout << "ok!\n";
/////////////////////////////////////////////////
string path = "/pytorch-deployment/assets/3.jpg";
Mat img = imread(path), img_float;
cvtColor(img, img, CV_BGR2RGB); //PyTorch的tensor是RGB通道,OpenCV是BGR
bitwise_not(img, img); // 将图像黑白反相
vector<Mat> mv;
split(img, mv);
img = mv[1]; // 截取单通道
img.convertTo(img_float, CV_32F, 1.0 / 255); // 使图像归一化
resize(img_float, img_float , Size(28, 28)); // 将图像resize成LeNet网络的输入大小
auto img_tensor = torch::from_blob(img_float.data, {1, 28, 28, 1}, at::kFloat).permute({ 0,3,1,2 }); // change the channels
auto img_var = torch::autograd::make_variable(img_tensor, false);
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(img_var);
auto output = module.forward(inputs).toTensor(); // 模型前向传播得到结果
cout << output << endl;
auto index = output.argmax(1); // 概率最大的onehot编码
cout << "The predicted class is : " << index << endl;
}
编译代码
由于我是用的 docker 版本的 OpenCV,同时又要用到代码源文件和 libtorch 库,就用 docker 的 -v
将物理机的目录挂载进 docker ,并用下面命令进入 docker (不在 docker 下部署的话可以忽视这一步)
$ docker run -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=$DISPLAY -p 5000:5000 -p 8889:8888 -e GRANT_SUDO=yes --user root -v ~/kevin/code/pytorch/pytorch-model-deployment:/pytorch-deployment -v ~/kevin/code/pytorch/libtorch:/libtorch -it spmallick/opencv-docker:opencv-3.4.3 /bin/bash
编译官方选用 CMake 工具,因此要写 CMakeLists.txt 文件,用 Qt Creator 构建我没有成功,以后有空再康康(主要是有代码自动补全功能)
cmake_minimum_required(VERSION 2.8.12)
set(CMAKE_INCLUDE_CURRENT_DIR ON)
set(CMAKE_AUTOMOC ON)
set(CMAKE_PREFIX_PATH ${CMAKE_PREFIX_PATH} "/libtorch")
project(pytorch-deployment)
find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)
add_executable(${PROJECT_NAME} "main.cpp")
target_link_libraries(${PROJECT_NAME} "${TORCH_LIBRARIES}")
target_link_libraries(${PROJECT_NAME} "${OpenCV_LIBS}")
set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 14)
依旧创建一个 build
目录来装编译中间产物,在 build 目录下输入以下命令构建可执行文件
$ cmake ..
$ make -j
构建完成后直接运行即可,不过这里还需要加上模型的存放路径作为命令行参数
$ ./pytorch-deployment /pytorch-deployment/assets/traced_mnist_pytorch_model.pt
可以得到最终的输出,本次用了一张手写数字 3 进行测试,模型最终输出也是 3,证明部署成功了,其实不出错的话,最终的结果和 PyTorch 会是一样的。
reference
https://pytorch.org/tutorials/advanced/cpp_export.html