Ubuntu下C++调用pytorch训练好模型--利用libtorch
最近因项目需要，需C++调用pytorch模型，以下是一些学习心得，把过程记录下来，同时供大家参考。
搭建有以下几个流程：
-
官网下载库libtorch
-
pytorch模型转化
-
编写C++调用程序
一、官网下载库libtorch
下载libtorch地址(https://pytorch.org/get-started/locally/），选择对应版本（根据自身电脑配置），比如我是Ubuntu+CUDA 10.2
下载好之后，选择路径进行解压（注意自己下载的包），命令如下：
unzip libtorch-shared-with-deps-1.5.1.zip解压之后会生成libtorch目录，目录下有如下文件夹，这是官方已编译好的文件，无需我们再编译，到这一步libtorch库准备完毕。
二、pytorch模型转化
这部分目的是将pytorch训练好的模型转化成C++可调用形式，创建tmp.py，直接贴代码了，执行python tmp.py，执行后会生成对应.pt文件
import torch
# An instance of your model.
from model.ResNet_models import ResNet
import torchvision.transforms as transforms
model = ResNet()
model.load_state_dict(torch.load('./ckpt/Resnet_New0703/epoch=30.pth'))
model.cuda()
model.eval()
# An example input you would normally provide to your model's forward() method.
testsize = 224
example = torch.rand(1, 3, testsize, testsize)
example = example.cuda()
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("./model_epoch=30.pt")注：所有路径都是相对路径，根据自身情况修改
执行过程中可能遇到类似如下错误：
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
原因是，输入数据类型和权重数据类型要统一，在上述代码，example = example.cuda()和model.cuda()就是解决统一的方案之一
三、编写C++调用程序
本部分需要做的流程：
- 新建文件夹，保存项目用的，比如src
- 创建CMakeLists.txt
- 编写test.cpp
- 编译执行
构建的项目包含文件及目录：
3.1 搭建程序所需环境，创建CMakeLists.txt
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(src) #src代表项目名
SET(CMAKE_C_COMPILER g++)
add_definitions(--std=c++14) #注意自己的c++版本，之前写的--std=c++11 报错，根据电脑配置定
# 指定libTorch位置
set(Torch_DIR /data/data1/hhq/20200707_libtorch/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)
message(STATUS "Torch library status:")
message(STATUS " version: ${TORCH_VERSION}")
message(STATUS " libraries: ${TORCH_LIBS}")
message(STATUS " include path: ${TORCH_INCLUDE_DIRS}")
message(STATUS " torch lib : ${TORCH_LIBRARIES} ")
# 需确保安装好OpenCV
find_package(OpenCV REQUIRED)
message(STATUS "OpenCV library status:")
message(STATUS " version: ${OpenCV_VERSION}")
message(STATUS " libraries: ${OpenCV_LIBS}")
message(STATUS " include path: ${OpenCV_INCLUDE_DIRS}")
message(STATUS " torch lib : ${TORCH_LIBRARIES} ")
# 包含头文件include
# include_directories(${OpenCV_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS})
add_executable(src test.cpp)
target_link_libraries(src ${TORCH_LIBRARIES} ${OpenCV_LIBS})
set_property(TARGET src PROPERTY CXX_STANDARD 14) # CXX_STANDARD 编号与add_definitions(--std=c++14)对应3.2 编写test.cpp
#include
#include "torch/script.h"
#include "torch/torch.h"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgcodecs.hpp"
#include
#include
#include
#include
using namespace cv;
using namespace std;
torch::Tensor process( cv::Mat& image,torch::Device device,int img_size)
{
vector mean_ = {0.485, 0.456, 0.406};
vector std_ = {0.229, 0.224, 0.225};
cv::cvtColor(image, image, cv::COLOR_BGR2RGB);// bgr -> rgb
cv::Mat img_float;
// image.convertTo(img_float, CV_32F, 1.0 / 255);//归一化到[0,1]区间,
cv::resize(image, img_float, cv::Size(img_size, img_size));
std::vector dims = {1, img_size, img_size, 3};
torch::Tensor img_var = torch::from_blob(img_float.data, dims, torch::kByte).to(device);//将图像转化成张量
img_var = img_var.permute({0,3,1,2});//将张量的参数顺序转化为 torch输入的格式 1,3,224,224
img_var = img_var.toType(torch::kFloat);
img_var = img_var.div(255);
for (int i = 0; i < 3; i++) {
img_var[0][i] = img_var[0][i].sub_(mean_[i]).div_(std_[i]);
}
return img_var;
}
int main()
{
/* 配置参数 */
char path[] = "../data/v1_20200713090006.jpg";
int img_size = 224;
torch::DeviceType device_type;
device_type = torch::kCUDA;
torch::Device device(device_type);
std::cout<<"cudu support:"<< (torch::cuda::is_available()?"ture":"false")<elements()[1].toTensor();
std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
std::cout << "Processing time = " << (std::chrono::duration_cast(t2 - t1).count())/1000000.0 << " sec" < 防坑，这有个Libtorch踩坑实录（https://www.jianshu.com/p/186bcdfe9492）：
- 如果模型返回值有多个，用toTuple()获取，不是直接toTensor()；具体看以上源码
3.3 编译执行
此部分搭建程序所需环境，在src文件下，新建build文件夹，并进入文件夹中，依次执行以下命令：
mkdir build
cd build
cmake ..
注：特别要注意CMakeLists.txt文件里边libtorch路径（第一步，解压后存放的路径），以及c++版本问题，这些都可能导致出错
执行cmake .. 成功后大致会出现以下信息
接下来，编译cpp文件，还是在build文件加下执行
make编译成功后，信息如下；注意每次修改test.cpp文件之后需要重新编译（make）
最后执行./src，查看结果
./src
参考：
Libtorch：pytorch分类和语义分割模型在C++工程上的应用
使用libtorch读取预训练权重，完成语义分割
Libtorch踩坑实录