论坛 › 深度学习 › Pytorch深度学习模型写入供C++调用

GTX_AI 发表于 2019-5-8 08:23:20

Pytorch深度学习模型写入供C++调用

Pytorch深度学习模型写入供C++调用import torch
from Models.MobileNetv2 import mobilenetv2

model = mobildnetv2(pretrained)
example = torch.rand(1, 3, 224, 224).cuda() # 注意，我这里导出的是CUDA版的模型，因为我的模型是在GPU中进行训练的
model = model.eval()

traced_script_module = torch.jit.trace(model, example)
output = traced_script_module(torch.ones(1,3,224,224).cuda())
traced_script_module.save('mobilenetv2-trace.pt')
print(output)

参考：
【1】https://zhuanlan.zhihu.com/p/52154049
【2】https://blog.csdn.net/IAMoldpan/article/details/85057238
【3】https://blog.csdn.net/IAMoldpan/article/details/86604302

GTX_AI 发表于 2019-5-8 23:11:17

import torch
import XXnet_path
net = XXnet()
net.cuda(0)      # GPU
net.load_state_dict(torch.load('xx.pth')['state_dict'])
net.eval()
example = torch.rand(1,3,256,256).type(torch.FloatTensor).cuda(0)
model = torch.jit.trace(net, example)
model.save( 'xx.pt' )

GTX_AI 发表于 2019-5-8 23:11:50

import torch
import XXnet_path
net = XXnet()
net.load_state_dict(torch.load('xx.pth')['state_dict'])
net.eval()
example = torch.rand(1,3,256,256)
model = torch.jit.trace(net, example)      # CPU
model.save( 'xx.pt' )

查看完整版本: Pytorch深度学习模型写入供C++调用