Torchscript格式模型 python demo

背景

• 前提：Pytorch导出Torchscript格式后不可逆；只能使用libtorch或者python的jit.load读取。
• 问题：队友只发了我模型和exe可执行文件没有代码，这样就不能愉快的玩耍了，百度了一圈网上python的demo不多，这里简单的总结一下。

查看

首先查看pt模型是否为Torchscript格式，如下图在netron中查看格式。

读取

import torch
model = torch.jit.load('model.pt')
#cpu
#model = torch.jit.load('model_cpu.pt',map_location=torch.device('cpu'))
#gpu 改变device(0)的参数即可，但需要对应模型转换时的配置gpu，我这里只能使用cpu。
#print(model)

推理测试

#手动构建一个0向量
img = torch.zeros(1,3,416,416)
out = model(img)[0]
print(out.size())

输出：
torch.Size([1, 10647, 8])

表示输出正常

后处理部分每个模型不同这里不细写，只要参照原本pytorch后处理方式即可。