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_scri
  • 4
    点赞
  • 67
    收藏
    觉得还不错? 一键收藏
  • 2
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值