Pytorch模型转Torchscript模型

最新推荐文章于 2025-01-27 09:09:06 发布

落花逐流水

最新推荐文章于 2025-01-27 09:09:06 发布

阅读量3.4k

点赞数 2

分类专栏： pytorch 深度学习 pytorch实践

本文链接：https://blog.youkuaiyun.com/juluwangriyue/article/details/108635280

版权

pytorch实践同时被 3 个专栏收录

136 篇文章 ¥29.90 ¥99.00

订阅专栏

pytorch

94 篇文章

订阅专栏

深度学习

19 篇文章

订阅专栏

本文介绍了将PyTorch模型转换为TorchScript的两种方法：追踪转换和注释转换。追踪转换简单但输入尺寸固定，而注释转换适用于包含特定控制流的模型。转换后的TorchScript模型可以脱离Python，在C++环境中运行。文章还提供了序列化Script模块到文件的步骤，并给出两个实用的转换函数示例。

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

TorchScript，它是PyTorch模型（子类nn.Module）的中间表示，可以在高性能环境（例如C ++）中运行。

转换的方法有两种，一种是通过追踪转换另一种是通过注释转换。

1、追踪转换

常用的是追踪转换，但是这种方法有一个缺点，就是输入尺寸固定。

官网给出的追踪的例子：

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

The traced ScriptModule can now be evaluated identically to a regular PyTorch module:

了解本专栏