第3章模型导出基础

模型导出，说白了就是把训练好的 PyTorch 模型“包装”成部署环境能直接吃的东西。很多同学在训练阶段跑得飞起，一到导出就踩坑——我当年也这样。这一章咱们把 PyTorch 保存与加载、ONNX 简介、TorchScript 入门一次性捋清楚。

3.1 PyTorch 模型保存与加载

我个人习惯把模型保存分成两种场景：训练中间存档 和 最终部署产物。别小看这一步，我曾经因为搞混了 torch.save 的两种方式，在项目交付前浪费了一整天。

核心 API 就两个：

torch.save(model.state_dict(), 'model.pth') —— 只存参数，轻量、推荐。
torch.save(model, 'model_full.pth') —— 存整个模型（包括结构），但容易出兼容问题。

你想想看，如果只保存 state_dict，加载的时候必须提前定义好模型类。我一般在项目里会单独写一个 model.py，把模型定义放进去，这样加载时 model.load_state_dict(torch.load('model.pth')) 就完事了。

# 保存
torch.save(model.state_dict(), 'resnet50_deploy.pth')

# 加载（需要先实例化模型）
model = ResNet50()
model.load_state_dict(torch.load('resnet50_deploy.pth'))
model.eval()

我曾经踩过的坑： 用 torch.save(model, ...) 保存完整模型后，换了 PyTorch 小版本就报错。所以现在我只存 state_dict，结构代码单独管理。

3.2 ONNX 简介

ONNX 是什么？说白了，它就是一个模型“通用语言”。你拿 PyTorch 训练好的模型，转成 ONNX 之后，就能在 TensorRT、ONNX Runtime、OpenVINO 上跑。我刚开始接触时觉得多此一举，直到有一次需要在 C++ 环境里做推理……嗯，ONNX 真香。

转换的核心是 torch.onnx.export。你需要给一个 dummy input，让 PyTorch 走一遍前向，然后记录计算图。

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'},
                  'output': {0: 'batch_size'}}
)

注意 dynamic_axes 这个参数。我建议你一开始就加上，否则导出的 ONNX 只能跑固定 batch size，部署时非常被动。

ONNX 模型可以用 onnxruntime 直接验证：ort.InferenceSession('model.onnx')。我每次导出后都会跑一遍，确保精度对齐。

3.3 TorchScript 入门

TorchScript 是 PyTorch 自家的“静态图”方案。它有两种模式：trace（追踪） 和 script（脚本）。我个人更常用 trace，因为它简单——给个输入，跑一遍，图就出来了。

# trace 模式
traced_model = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
traced_model.save('model_traced.pt')

# script 模式（适合有控制流的模型）
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

为什么需要 TorchScript？你想想看，Python 解释器在线上环境里有时候是累赘。TorchScript 可以脱离 Python 运行，还能被 C++ 直接调用。我之前在嵌入式设备上部署，就是用 TorchScript 把模型打包，然后通过 LibTorch 加载，延迟直接降了 30%。

三种导出方式对比：

方式	便携性	跨框架	推理性能	我的推荐场景
state_dict	⭐⭐⭐	仅 PyTorch	原生	训练存档、调试
ONNX	⭐⭐⭐⭐⭐	强	优（经优化器）	生产部署、多后端
TorchScript	⭐⭐⭐⭐	仅 PyTorch/C++	优秀	纯 PyTorch 管线、C++ 集成

注意： TorchScript 的 trace 模式无法捕捉动态控制流（比如 if-else 依赖数据）。如果你的模型里有循环或条件分支，请用 torch.jit.script。我曾经在 trace 一个带 attention mask 的模型时，结果完全不对……后来换成 script 才搞定。

好了，这一章的基础知识就这些。记住：导出不是终点，而是部署的起点。下一章我们会拿 ONNX 做更深入的实战，到时候你会感谢今天打下的基础。

第3章 模型导出基础

3.1 PyTorch 模型保存与加载

3.2 ONNX 简介

3.3 TorchScript 入门

第3章模型导出基础