PyTorch JIT 本身不直接加速推理，核心价值是固化模型结构与支持跨平台部署；真正加速依赖后续图优化和后端适配，需手动启用如 fold_conv_bn 等 pass。
PyTorch JIT 能不能直接加速推理？别盲目上 torch.jit.script
大多数场景下，torch.jit.script 不会自动提速，甚至可能变慢。它核心价值是固化模型结构、剥离 Python 控制流、生成可序列化/跨进程部署的 ScriptModule，而非“编译优化”。真正带来推理加速的，通常是后续的图级优化（如融合、常量折叠）和后端适配（如 TorchScript + CPU backend 的向量化），但这些不依赖你手动调用 JIT —— torch.jit.trace 或 torch.jit.script 之后再启用 torch._C._jit_pass_fold_conv_bn 等 pass 才算真正介入。
实操建议：
- 先用
torch.jit.trace尝试，尤其适合输入 shape 固定、无 if/for 动态逻辑的模型（如典型 CNN 推理） - 遇到
TracingFailed或输出异常，再切到torch.jit.script，但要检查所有分支是否都被覆盖，否则 runtime 报错是常态 - 务必在 trace/script 后调用
.eval()和.to(device)，否则 JIT 模块仍会走 Python fallback 路径
trace 时输入 tensor 必须满足什么条件？
torch.jit.trace 本质是记录一次前向执行路径，所以输入必须能完整驱动模型所有分支，且 shape、dtype、device 必须与实际推理一致。常见翻车点：
- 输入是
None或含可选参数（如mask=None）→ trace 时该分支被跳过，后续推理传入非 None 会 crash - batch size 用 1 trace，但线上跑 batch=16 → 图中所有 tensor shape 被硬编码为 [1, ...]，运行时报 dimension mismatch
- CPU 上 trace，GPU 上运行 → JIT 模块里权重仍在 CPU，触发隐式 copy，性能暴跌
正确做法：用真实推理环境下的最小 batch 输入（如 torch.randn(1, 3, 224, 224).to('cuda')）做 trace，并确保模型已 .eval()。
立即学习“Python免费学习笔记（深入）”；
script 编译失败常见报错怎么快速定位？
最典型的错误是 NotSupportedError: with statement is not supported 或 UnsupportedNodeError: 'Dict' object is not subscriptable，说明 Python 动态特性超出了 TorchScript 类型系统能力。
- 避免在 forward 中直接用
dict.keys()、list.append()、**kwargs解包 - 把 dict 操作提前转成 tuple/list，或用
@torch.jit.ignore标记纯 Python 辅助函数（但注意：被 ignore 的部分无法被优化） - 用
torch.jit.script前先运行torch.jit.fuser("fuser2")（新版默认启用），有时能绕过某些旧 fuser 的限制 - 加
torch.jit.set_script_logging(True)查看具体哪行被拒绝
保存和加载后的模块还能不能改输入 shape？
可以，但仅限于支持动态维度的模型。JIT 模块本身不绑定 shape，真正限制来自 traced graph 中的常量节点或算子约束（如 nn.AdaptiveAvgPool2d(1) 输出固定为 1×1）。关键看 trace 时是否用了 symbolic shapes：
- PyTorch 1.10+ 支持
torch.jit.trace的example_inputs传入带torch.export.Dim的 symbolic shape，但目前仍属 experimental - 更稳妥的方式：用
torch.jit.freeze+torch.jit.optimize_for_inference组合，后者会启用更多图优化（如算子融合、内存复用），对 batch size 变化容忍度更高 - 加载后调用
.cuda()或.to(device)是安全的；但.train()会破坏优化效果，且可能引发未定义行为
真正容易被忽略的是：JIT 优化效果高度依赖模型结构和硬件后端。一个在 V100 上提升 20% 的 traced 模块，在 A10 上可能只有 5%，甚至因 kernel 选择不佳而倒退。上线前务必在目标设备上实测吞吐和延迟，别只信 benchmark 数值。
