在C++平台部署pytorch模型的两种方法

1、准备工作

1、模型转换代码：

import torch
from model.FCN_Aux import FCN_AUX
model = FCN_AUX(2)
model_weight_path = './model/better.pth'#训练完模型位置
model.load_state_dict(torch.load(model_weight_path))
model.cuda()
model.eval()
example = torch.rand(1, 3, 224, 224).cuda()
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("./model/model__gpu_50.pt")#保存模型位置

2、下载对应版本pytorch
pytorch官网连接

将pytorhc加入环境变量

2、通过cmake构建配置

pytorch官网示例连接
1、文件结构

2、cmakelists
这里我还用到了opencv

cmake_minimum_required(VERSION 3.16 FATAL_ERROR)
project(test)
 
set(CMAKE_PREFIX_PATH "E:\\lib\\libtorch_gpu\\libtorch\\share\\cmake\\Torch" )
set(OpenCV_DIR "E:\\lib\\opencv_4.1.0\\opencv\\build")
find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)
 
if(NOT Torch_FOUND)
    message(FATAL_ERROR "Pytorch Not Found!")
endif(NOT Torch_FOUND)
 
message(STATUS "Pytorch status:")
message(STATUS "    libraries: ${TORCH_LIBRARIES}")
 
message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")
 
 
add_executable(test test.cpp)
target_link_libraries(test ${TORCH_LIBRARIES} ${OpenCV_LIBS})
set_property(TARGET test PROPERTY CXX_STANDARD 11)

3、test.cpp
test.cpp里就看项目需要

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include

cv::Mat find_pin(cv::Mat img);
int main() {
	clock_t t1, t2;
	t1 = clock();
	cv::Mat src = cv::imread("0.jpg");
	cv::Mat label = find_pin(src);
	cv::imwrite("label.jpg", label);
	t2 = clock();
	std::cout << t2 - t1 << std::endl;
	system("pause");
}

cv::Mat find_pin(cv::Mat img)
{
	torch::Device device(torch::kCUDA);
	torch::NoGradGuard no_grad;
	torch::jit::script::Module module = torch::jit::load("model_pin_gpu_50.pt");
	module.to(device);

	img.convertTo(img, CV_32F, 1.0 / 255.0);
	auto img_tensor = torch::from_blob(img.data, { 1,1571,1920,3 });
	img_tensor = img_tensor.permute({ 0,3,1,2 });
	img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229);
	img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
	img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);

	img_tensor = img_tensor.to(device);

	torch::Tensor output = module.forward({ img_tensor }).toTensor();

	torch::Tensor out_tensor = output.argmax(1);
	out_tensor = out_tensor.squeeze();

	out_tensor = out_tensor.mul(126).to(torch::kU8);
	out_tensor = out_tensor.to(torch::kCPU);

	cv::Mat resultImg(1571, 1920, CV_8U);
	std::memcpy((void *)resultImg.data, out_tensor.data_ptr(), sizeof(torch::kU8) * out_tensor.numel());

	return resultImg;
}

代码段参考了其他博客，连接如下：
深度学习工程实践 5. libtorch+opencv使用复杂的训练分割模型(脸部，头发实现分割)
3、建立工程
打开命令提示符，并cd到build文件夹下。
输入`

cmake -DCMAKE_GENERATOR_PLATFORM=x64 ..
cmake -DCMAKE_PREFIX_PATH=E:/lib/libtorch_gpu/libtorch ..
cmake --build . --config Release

4、build文件夹

打开test.sln,移除ALL_BUILD和ZERO_CHECK,改为release.
把模型和图片放到指定位置运行即可，QT也可以通过类似方法实现。

3、通过配置项目属性表

参考：Libtorch：pytorch分类和语义分割模型在C++工程上的应用
添加包含目录和库目录

在附加依赖项中加入

c10.lib
c10_cuda.lib
caffe2_module_test_dynamic.lib
torch.lib

将SDL检查改为否

将libtorch文件夹下的lib目录中的文件复制到bin目录下

运行即可。