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;

你可能感兴趣的:(c++,pytorch,开发语言)