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

GTX_AI

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

1. import torch
   
2. from Models.MobileNetv2 import mobilenetv2
   
3. 
   
4. model = mobildnetv2(pretrained)
   
5. example = torch.rand(1, 3, 224, 224).cuda() # 注意，我这里导出的是CUDA版的模型，因为我的模型是在GPU中进行训练的
   
6. model = model.eval()
   
7. 
   
8. traced_script_module = torch.jit.trace(model, example)
   
9. output = traced_script_module(torch.ones(1,3,224,224).cuda())
   
10. traced_script_module.save('mobilenetv2-trace.pt')
    
11. 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

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

复制代码

GTX_AI

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

复制代码