torchscript保存模型和使用模型

torchscript作为pytorch的模型保留方式

这里主要介绍3种类型：

1.单变量输入

x = torch.randn(1, 3, 256, 256, requires_grad=True).cuda()
script_module = torch.jit.trace(model,x,strict=False)
torch.jit.save(script_module, "model.pt")

2.多个变量作为输入

x1 = torch.randn(1, 3, 256, 256, requires_grad=True).cuda()
x2 = torch.randn(1, 3, 256, 256, requires_grad=True).cuda()
x3 = torch.randn(1, 3, 256, 256, requires_grad=True).cuda()
script_module = torch.jit.trace(model,(x1,x2,x3),strict=False)
torch.jit.save(script_module, "model.pt")

 注意：多个输入的时候，上述的x1,x2,x3是可以任意的，只要使用括号括起来就行。

3.分布式训练的模型保存成torchscript

x = torch.randn(1, 3, 256, 256, requires_grad=True).cuda()
script_module = torch.jit.trace(model.module,x,strict=False)
torch.jit.save(script_module, "model.pt")

没看错，就是加个module就行，这样才是模型，不然是分布式的数据结构 

至此，常用的保存为torchscript就介绍到这里

至于加载模型也很简单

model=torch.jit.load("model.pt")