在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