C++部署模型

在人工智能领域，Python受到学术界的追捧，模型训练比模型部署性能更加重要。然而在实际终端部署方面，低延迟、可移植性和可适用性的需求使得Python成为一个比较差的语言。相反，C++凭借其可移植性、可适用性以及运算速度快等优势，更适合终端部署网络模型。下面我将以我做的部署ReID模型为例子，简要介绍如何利用Libtorch（or: Pytorch for C++ API）实现C++部署Pytorch训练的模型。

环境搭建

  • VS2019
  • OpenCV
  • 4.5.5 + release
  • Libtorch
  • 1.8.1 + release +cpu
  • Cmake

为了使得推理速度更快，OpenCV和Libtorch都推荐使用release版本。下面是几点环境搭建的注意点，很重要！！！

  1. Pytorch和Litorch的版本最好一致
  • 否则就会出现模型加载不了的问题
  1. OpenCV release版本链接库是不带d的
  • 比如我的是，release版本的库文件名为：opencv_world455.lib，debug版本的库文件名为：opencv_world455d.lib
  1. CUDA版本也要和Pytorch一致
  • 我电脑没有独显，就直接用的cpu版本的libtorch

部署步骤

Libtorch部署Pytorch模型主要包括以下流程：

  1. Python端导出模型
  2. C++端加载模型
    3.C++执行Script

Python端导出模型

Pytorch训练保存得到的pt模型（pt, pth, pkl只是文件名不同，无本质区别）和C++端加载的pt模型不是同一个东西！实现模型由Python向C++转换，需要依赖torchscript，其保存的后缀名常常简化为.pt（.torchscript.pt)。torchscript能够被torchscript编译器推理、编译以及序列化。要想实现Pytorch模型转换为torchscript，有tracing和annotation两种方式。下面将将简要介绍tracing方式。

tracing

就是向pytorch模型注入一个简单输入，使用torch.jit.trace对这个流进行跟踪，从而得到torchscript模型。下面是我写的转换pytorch模型的一个例子。

import torch
import torchvision

model = net(class_num=751)
device = torch.device('cpu')
model.load_state_dict(toch.load('best.pt', map_location=device))

x=Variable(torch.FloatTensor(8, 3, 256, 128))
traced_script_model = torch.jit.trace(model, x)
traces_script_model.save('best.torchscript.pt')

torch.jit.trace对这个流进行跟踪, 使用 .save方法序列化模型，以便不依赖python。
注意：

  • 输入要和网络输入大小一致
  • pytorch加载模型的方式要和其保存模型相搭配

C++端加载模型

CMake构建Libtorch工程

先用CMake构建好Libtorch工程。项目文件的组织形式如下：

  • reid
  • include
  • src
  • bin
  • weights
  • data
  • CMakeLists.txt
    下面是CMakeLists文本文件的内容：
cmake_minimum_required (VERSION 3.12)
project (reid_demo_libtorch)

set (TORCH_DIR "E:\\libtorch1.8.1\\libtorch\\share\\cmake\\Torch")
set (OpenCV_DIR "E:\\opencv\\opencv\\build")
find_package (Torch REQUIRED)
find_package (OpenCV REQUIRED)

if (NOT Torch_FOUND)
	message (FATAL_ERROR "Libtorch Not Found!")
endif (NOT Torch_FOUND)
if (NOT OpenCV_FOUND)
	message (FATAL_ERROR "OPenCV Not Found!")
endif (NOT OpenCV_FOUND)

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

aux_source_directory (src SRC_LIST)
include_directories (include)

add_executable (main ${SRC_LIST}) 
target_link_libraries (
	main
	${OpenCV_LIBS}
	${TORCH_LIBRARIES}
)

上面文件的内容主要是找到OpenCV和Libtorch库的位置，输出相关信息，再绑定库。具体的CMake语法可以参考CMake - Linux基础教程

  1. cd reid
  2. mkdir build 建立一个空白的build目录，
  • CMake构建的项目工程会存放到build目录下
  1. cd build
  3. cmake ..cmake -DCMAKE_PREFIX_PATH=E:\opencv\opencv\build\x64\vc15\lib;E:\libtorch1.8.1\libtorch -DCMAKE_BUILD_TYPE=Release -G "Visual Studio 16 2019" ..
    • 上面的位置信息根据自己的做相应更改
    • G "Visual Studio 16 2019"是vs2019的编译工具
    • 加上编译类型，Release
C++加载模型

加载序列化的torchscript模型的流程如下：

std::string model_path = "../weights/best.torchscript.pt";
torch::jit::script::Module model;
try {
	model = torch::jit::load(model_path);
}
catch (const c10::Error& error) {
	std::cerr << "error loading model!" << std::endl;
	std::exit(EXIT_FAILURE);
}
model.to(torch::kCPU);
model.eval（);

torch::jit::load加载torchscript模型，并对模型反序列化，返回torch::jit::script::Module模型对象。

C++执行Script

下面将介绍C++读取输入数据，利用模型做前向推理的过程。

std::vector<torch::jit::IValue> inputs;
auto input = PreProcessReID(img);
inputs.emplace_back(input);

torch::Tensor outputs = model.forward(inputs).toTensor();
torch::Tensor feature = outputs.contiguous().view({1, -1});