torch.jit.script更适合含控制流、动态结构或自定义python逻辑的模型，torch.jit.trace仅记录单次执行路径，对条件分支和输入shape变化敏感；标准cnn/transformer且输入固定时trace更轻量高效。
torch.jit.script 和 torch.jit.trace 哪个更适合你的模型？
torch.jit.script 适合含控制流（如 if、for）、动态结构或自定义 Python 逻辑的模型；torch.jit.trace 只记录一次前向执行路径，对条件分支、输入 shape 变化敏感，容易“漏逻辑”。
- 如果模型里有
if x.size(0) > 1:或循环展开 token，必须用torch.jit.script - 如果只是标准 CNN/Transformer 且输入 shape 固定，
torch.jit.trace更快、更轻量 -
torch.jit.script会校验类型，遇到list.append()或未标注的可选参数可能报RuntimeError: Unable to cast Python instance - trace 时传入的示例输入必须覆盖所有可能分支，否则导出后运行会出错（比如某分支没执行过，trace 后直接消失）
导出前必须处理的三类 Python 依赖
TorchScript 不支持任意 Python 运行时行为，常见卡点：
-
print()、logging.info()、assert：全部删掉或用torch._assert()替代（注意后者只在 debug 模式生效） - 字典推导、
**kwargs解包、getattr(model, name)：TorchScript 无法静态解析，需改写为显式属性访问或预定义分支 - NumPy 调用（如
np.array()、np.where()）：必须替换为等价的torch.tensor()或torch.where()，否则script()直接失败
例如这段代码无法 script：
def forward(self, x):
if x.shape[0] == 1:
return self.small_head(x)
else:
return self.large_head(x)
要改成：
def forward(self, x):
batch_size = x.shape[0]
torch._assert(batch_size > 0, "batch size must be positive")
if batch_size == 1:
return self.small_head(x)
else:
return self.large_head(x)
保存和加载 .pt 文件时的兼容性陷阱
torch.jit.save() 输出的是二进制序列化文件，不是纯权重，它包含模型结构、常量、优化后的图——但不包含原始 Python 类定义。
- 加载时不需要原始代码，但若用了
torch.jit.script且类中引用了外部函数（如自定义激活函数），该函数也必须被script过，否则torch.jit.load()报Unknown builtin op - 不同 PyTorch 版本间 .pt 文件不保证向前/向后兼容：1.12 导出的模型在 2.0 上可能加载失败，错误信息类似
version mismatch: expected X, got Y - 若需跨环境部署（如 C++ backend），务必用
model.save("model.pt")而非torch.jit.save(model, ...)，前者更稳定
验证导出模型是否等价的最简方法
别只看能否加载，重点比输出数值和梯度：
- 用同一组输入跑原模型和导出模型，检查
torch.allclose(out1, out2, atol=1e-5) - 若模型参与训练，还要测试
out.sum().backward()后，两者的model.weight.grad是否一致（注意 trace 模型默认不带 backward 图，需额外调用model = torch.jit.script(model)并确保 forward 中无 in-place 操作） - 在 CPU 上验证完再移至 GPU：有些算子（如
torch.nn.functional.scaled_dot_product_attention）在 trace 时可能因 device 不匹配 silently fallback，导致结果偏差
真正麻烦的往往不是导出动作本身，而是模型里那些“看起来没问题”的 Python 惯用写法——它们在 TorchScript 下根本不会报错，只是静默失效。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
