torchscript保存模型和使用模型

最新推荐文章于 2025-03-06 18:04:30 发布

喝粥也会胖的唐僧

最新推荐文章于 2025-03-06 18:04:30 发布

阅读量6.3k

点赞数 1

分类专栏：深度学习

本文链接：https://blog.youkuaiyun.com/zhou_438/article/details/115186842

版权

深度学习专栏收录该内容

78 篇文章

订阅专栏

本文介绍了如何使用TorchScript将PyTorch模型保存为三种类型：单变量输入、多个变量输入和分布式训练模型。通过torch.jit.trace方法结合cuda设备，对模型进行trace并保存为model.pt。加载模型则通过torch.jit.load实现。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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就介绍到这里

至于加载模型也很简单