C++调用已训练好的pytorch模型

使用C++调用pytorch模型,即将模型部署到生产环境中,这里先提供一个pytorch的官方教程

首先需要下载并配置好torch script环境,网上已有相关教程。然后在python中将模型转换为torch script,我使用了annotation的方式。

net = Net().to(device)
net.load_state_dict(torch.load('unified_model.pth', 'cpu'))

# convert pytorch model to Torch Script via annotation(another is tracing)
sm = torch.jit.script(net)
sm.save("converted_model.pt")

# test output
f = torch.tensor([[1.,2.,3.,4.,5.]])
v = torch.tensor([[1.,2.,3.,4.,5.,1.,2.,3.,4.,5.]])
t = torch.tensor([[1.,2.,3.,4.,5.,1.,2.,3.,4.,5.,1.,2.,3.,4.,5.,1.,2.,3.,4.]])

out = net(f,v,t)

随后在C++中加载保存好的模型,并使用相同的样例测试模型输出。

torch::jit::script::Module model;
model = torch::jit::load("D:\\pythoncode\\converted_model.pt"); // load the model

float f[] = { 1,2,3,4,5 };
float v[] = {1,2,3,4,5,1,2,3,4,5};
float t[] = { 1,2,3,4,5,1,2,3,4,5,1,2,3,4,5,1,2,3,4 };
torch::Tensor f1 = torch::from_blob(f, { 1,5 });
torch::Tensor v1 = torch::from_blob(v, { 1,10 });
torch::Tensor t1 = torch::from_blob(t, { 1,19 });

torch::Tensor output = model.forward({ f1,v1,t1 }).toTensor();
std::cout << output << std::endl;
  • 1
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值