赞
踩
最近因项目需要,需C++调用pytorch模型,以下是一些学习心得,把过程记录下来,同时供大家参考。
下载libtorch地址(https://pytorch.org/get-started/locally/),选择对应版本(根据自身电脑配置),比如我是Ubuntu+CUDA 10.2
下载好之后,选择路径进行解压(注意自己下载的包),命令如下:
unzip libtorch-shared-with-deps-1.5.1.zip
解压之后会生成libtorch目录,目录下有如下文件夹,这是官方已编译好的文件,无需我们再编译,到这一步libtorch库准备完毕。
这部分目的是将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
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。