在人工智能领域,Python受到学术界的追捧,模型训练比模型部署性能更加重要。然而在实际终端部署方面,低延迟、可移植性和可适用性的需求使得Python成为一个比较差的语言。相反,C++凭借其可移植性、可适用性以及运算速度快等优势,更适合终端部署网络模型。下面我将以我做的部署ReID模型为例子,简要介绍如何利用Libtorch(or: Pytorch for C++ API)实现C++部署Pytorch训练的模型。
为了使得推理速度更快,OpenCV和Libtorch都推荐使用release版本。下面是几点环境搭建的注意点,很重要!!!
d
的
opencv_world455.lib
,debug版本的库文件名为:opencv_world455d.lib
Libtorch部署Pytorch模型主要包括以下流程:
Pytorch训练保存得到的pt模型(pt, pth, pkl只是文件名不同,无本质区别)和C++端加载的pt模型不是同一个东西!实现模型由Python向C++转换,需要依赖torchscript
,其保存的后缀名常常简化为.pt
(.torchscript.pt
)。torchscript
能够被torchscript
编译器推理、编译以及序列化。要想实现Pytorch模型转换为torchscript
,有tracing
和annotation
两种方式。下面将将简要介绍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。
注意:
先用CMake构建好Libtorch工程。项目文件的组织形式如下:
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基础教程
cd reid
mkdir build
建立一个空白的build
目录,
build
目录下cd build
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
加载序列化的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++读取输入数据,利用模型做前向推理的过程。
std::vector inputs;
auto input = PreProcessReID(img);
inputs.emplace_back(input);
torch::Tensor outputs = model.forward(inputs).toTensor();
torch::Tensor feature = outputs.contiguous().view({1, -1});