PyTorch导出TorchScript完整教程

本文深入解析了PyTorch模型导出为TorchScript的两种核心方法——tracing与scripting的本质区别、适用场景及典型陷阱：tracing仅记录单次前向执行路径，对输入shape和控制流极度敏感，易因分支未覆盖或device/dtype不一致导致静默错误；scripting则通过静态分析支持条件逻辑和循环，但严格限制Python语法与类型表达，常见于含训练/评估切换或动态层数的模型。文章不仅指出混合使用策略的可行性，更强调导出后验证的关键性——从输入预处理、设备状态、随机性控制到图结构检查，每一步疏漏都可能引发线上推理异常，真正决定落地成败的，往往不是导出那行代码，而是紧随其后的严谨验证。

tracing 和 scripting 的核心区别在哪

tracing 是“记录一次前向执行路径”，只捕获你给的示例输入走过的计算图；scripting 是“静态分析 Python 代码”，能处理控制流（比如 if、for），但要求所有逻辑可被 JIT 编译器解析。

常见错误现象：torch.jit.trace 导出后模型在推理时输出全零、形状错乱，或遇到 RuntimeError: input is not a tensor —— 大概率是 trace 没覆盖到分支逻辑，比如训练/评估模式切换、条件 dropout、动态 shape 处理。

• 用 tracing：适合结构固定、输入 shape 稳定、不含复杂 Python 控制流的模型（如标准 ResNet、CNN）
• 用 scripting：适合含 if self.training:、for i in range(self.num_layers): 或自定义 forward 中有明确逻辑分支的模型
• 混合用法可行：先 torch.jit.script 主干，再对部分子模块 torch.jit.trace（需确保 traced 模块不依赖外部变量）

trace 时输入必须和部署时完全一致

trace 不会泛化输入 shape 或 dtype，它记住的是“这次传进来的 input 长什么样”。后续用不同 shape 的 tensor 推理，轻则报错，重则静默出错（比如 padding 错位）。

典型坑：torch.jit.trace(model, torch.randn(1, 3, 224, 224)) 导出后，拿 torch.randn(4, 3, 224, 224) 去跑，可能 crash 或结果异常 —— 因为某些算子（如 BatchNorm2d）在 trace 时固化了 batch size 相关的统计量或 buffer 形状。

• trace 前务必调用 model.eval()，否则 dropout、BN 行为不可控
• 输入 tensor 必须在相同 device 上（比如都用 .cuda()），否则 trace 出的图里 device 信息错乱
• 如果模型接受多输入（如图像 + 文本 embedding），必须传 tuple 或 dict，且 key 名/顺序要和 forward 签名严格一致

scripting 报错常见原因和绕过方法

最常卡在 Not supported: augmented assignment 或 Unsupported primitive type: typing.Union —— 这些不是 bug，是 TorchScript 类型系统故意拦住不确定行为。

比如模型里写了 self.count += 1，TorchScript 不支持原地修改属性；又或者用了 Optional[Tensor]，但没显式处理 None 分支。

• 把可变状态（计数器、缓存）移到 forward 外，或改用 torch.tensor 作为参数传入
• 所有 Optional 参数必须显式判断：不能只写 if x is not None:，得补上 else: 分支，哪怕只是 return torch.zeros(...)
• 避免用 Python list/dict 存 tensor，改用 torch.nn.ModuleList 或 torch.stack；list comprehension 要重写为显式循环
• 第三方库函数（如 sklearn、cv2）无法 script，得提前移出 forward

导出后验证比导出本身更重要

很多人导出成功就以为完事了，但 torch.jit.load 加载后直接 run，和原始模型输出差 1e-3 以上很常见——这往往不是精度问题，而是 trace/script 漏掉了某个隐式调用（比如 model.to(device) 后没同步 buffer，或 torch.no_grad() 没包住）。

验证时最容易忽略的是：是否用了和原始模型完全一致的预处理、device、dtype、以及 torch.set_grad_enabled(False) 状态。

• 必须对比同一组输入下，原始模型和 scripted_model 的输出 tensor，用 torch.allclose(out1, out2, atol=1e-5) 而不是 ==
• 如果模型含随机性（如未设 seed 的 dropout），先全局禁用：torch.manual_seed(0); torch.backends.cudnn.deterministic = True
• 检查 scripted_model.graph 是否包含可疑节点（如 prim::If 消失了，或出现 aten::python_call —— 说明某处 fallback 到了 Python，性能会崩）

真正麻烦的从来不是怎么导出，而是导出后那几行验证代码没写全。漏掉一个 .eval() 或一个 .cpu()，线上 inference 就可能返回 NaN。

好了，本文到此结束，带大家了解了《PyTorch导出TorchScript完整教程》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！