4. PyTorch 模型导出：torch.jit.trace 与 torch.jit.script 的区别与实战
模型训练好了，怎么把它搬到生产环境去？这是每个部署工程师都要面对的问题。PyTorch 提供了两种导出方式：torch.jit.trace 和 torch.jit.script。说白了，一个靠「看」来理解模型，一个靠「读」来理解模型。
我个人习惯把 trace 叫做「傻瓜式导出」，script 叫做「聪明式导出」。为什么这么叫？往下看你就明白了。
4.1 torch.jit.trace：跟踪执行路径
Trace 的原理很简单——给模型一个输入，它跑一遍，把每一步操作都记录下来。就像你开车走一遍路，它就把导航轨迹记下来了。
核心特点：
- 只记录实际执行到的代码路径
- 不支持动态控制流（if/for 依赖数据的情况）
- 导出速度快，生成的模型体积小
来看个例子：
import torch
class MyModel(torch.nn.Module):
def forward(self, x):
return x * 2 + 1
model = MyModel()
example_input = torch.randn(1, 3, 224, 224)
# 使用 trace 导出
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")嗯，这里要注意。Trace 有个大坑——它只记录你给的那个输入下的执行路径。我在项目中遇到过一个问题：模型里有个 if x.sum() > 0 的分支，测试时输入都是正数，结果导出后部署到线上，遇到负数输入就直接崩了。
避坑指南：我曾经因为 trace 导出导致线上服务挂了两个小时。原因是模型里有个动态的循环，trace 只记录了第一次迭代的路径。所以，如果你的模型有数据依赖的分支或循环，千万别用 trace。
4.2 torch.jit.script：编译计算图
Script 就不一样了。它直接读取你的 Python 代码，把整个控制流都编译成 TorchScript 的中间表示。你想想看，这就像把菜谱完整地翻译成另一种语言，而不是只看你做一次菜。
核心特点：
- 完整保留控制流逻辑
- 支持动态图结构
- 导出速度慢一些，但更安全
import torch
class MyModel(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x * 2
else:
return x * 0.5
model = MyModel()
# 使用 script 导出
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")Script 也不是万能的。它要求你的代码必须是「TorchScript 兼容」的。什么意思？就是你不能用那些 Python 特有的语法，比如 yield、with 语句，或者某些第三方库。
4.3 核心区别对比
| 对比维度 | torch.jit.trace | torch.jit.script |
|---|---|---|
| 工作原理 | 跟踪执行路径 | 编译源代码 |
| 控制流支持 | 不支持动态控制流 | 完整支持 |
| 导出速度 | 快 | 较慢 |
| 代码兼容性 | 几乎无限制 | 需 TorchScript 兼容 |
| 调试难度 | 低 | 中等 |
| 推荐场景 | 简单前馈网络 | 复杂控制流模型 |
4.4 实战：如何选择？
我一般遵循这个原则：
- 先试 trace——如果模型没有动态控制流，trace 是最快最稳的选择
- trace 报错再换 script——遇到
UnsupportedNodeError之类的错误，说明模型里有动态逻辑 - 混合使用——有时候可以把模型拆开，静态部分用 trace，动态部分用 script
小技巧：我个人习惯在开发阶段先用 script 导出，确保所有逻辑都被正确编译。等到模型稳定了，再考虑用 trace 优化性能。这样既保证了正确性，又兼顾了效率。
4.5 常见问题与解决方案
Q: trace 导出的模型在 C++ 里跑不起来？
A: 检查一下输入输出的 dtype 和 shape 是否一致。我曾经遇到过一个 bug，训练时用的是 float32，但 trace 时给了 float64 的输入，结果导出的模型在 C++ 端死活加载不了。
Q: script 导出时提示 Cannot insert a Tensor that requires grad？
A: 这说明你的模型里还有需要梯度的张量。导出前记得调用 model.eval() 和 torch.no_grad()。
Q: 模型里有自定义的 nn.Module 怎么办？
A: 只要你的自定义模块继承自 nn.Module，并且 forward 方法里只用了 PyTorch 支持的操作，script 和 trace 都能处理。如果用了第三方库，那就得考虑用 @torch.jit.ignore 装饰器跳过某些方法。
好了，这一章就到这里。记住一句话：trace 是快刀，script 是钝刀。快刀能切大部分菜，但遇到骨头还是得用钝刀。下一章我们聊聊 ONNX 导出，那又是另一个故事了。