最近因项目需要,需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.
t