【pytorch】Pytorch训练好的模型转torchscript

最新推荐文章于 2025-02-15 10:44:26 发布

一方热衷.

最新推荐文章于 2025-02-15 10:44:26 发布

阅读量1.8k

点赞数 2

文章标签：人工智能 python pytorch

本文链接：https://blog.csdn.net/qq_44992785/article/details/129346751

版权

文章介绍了如何将基于PyTorch的nn.Module模型转换为TorchScript，以便在高性能环境中运行，如C++。通过torch.jit.trace函数对模型进行追踪，然后保存转换后的模型。加载预训练权重，对模型进行评估，并使用固定输入维度的示例数据保存trace后的脚本模块。

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

TorchScript，它是PyTorch模型（子类nn.Module）的中间表示，可以在高性能环境（例如C ++）中运行。利用torch.jit.trace()函数

import torch
import torchvision
model=get_model #获得你的模型
model.load_state_dict(torch.load('./best.pth', map_location='cpu'))#加载你训练好的模型权重
model  = model.eval()#验证
dummy_input = torch.randn(1,3,299,299)#固定的输入维度
traced_script_module = torch.jit.trace(model, dummy_input)#
traced_script_module.save('./tt.pt')#保存路径

ps：如果你训练的时候是模型和参数保存在一起了直接model=torch.load('.pth')就可以了，上面代码演示的权重单独保存的记载方式