pt模型转torchscript模型

在GitHub上找了个图像处理的工程,它自带了个pt文件,

model =a.LResNet50E_IR_Occ(num_mask=226)
test_model=torch.load("model_p5_w1_9938_9470_6503.pth.tar")
print(test_model.keys())

打印是
dict_keys(['state_dict', 'classifier'])
首先,初识模型转换,加载它自带的这个pt文件的model有state_dict和classifer两个键,第一个问题:
如果想转换成torchscript是不是只把state_dict加载进本工程的网络结构就可?state_dict对应的值又是一个 <class 'collections.OrderedDict'> ,键是各个网络层的权重。
第二个问题:
通用的将pt模型转torchScript的步骤是什么?

img


我的这个步骤有没有问题?


通常普通的简单步骤可以看官网的实例,
一般来说你要先保证你的模型能够正常运行出结果,然后直接
ts = torch.jit.trace(model, input_, strict=False)
ts.save(path/to/torchscript)
比较麻烦的是这个_extra_files,看你自己的网络需要存哪些东西下来了,如果不需要的话(也就是网络有些参数你去不知道,不需要模型帮你保存,就忽略)