📚 课程目录第2章模型导出与格式转换

模型部署与推理优化案例 - 第2节

第 2 章：模型导出与格式转换——ONNX、TorchScript、SavedModel

模型训练好了，只是第一步。

说实话，在生产环境里，没人直接拿 PyTorch 或 TensorFlow 的训练权重去做推理。你想想看，训练框架太重了，几百兆的依赖，还有各种动态图的性能开销。所以，模型导出和格式转换，就是我们部署前必须过的一道坎。

我个人习惯把这一阶段叫做「模型标准化」。就是把训练好的模型，转成一种中间表示，让不同硬件、不同推理框架都能识别。今天咱们就聊聊三种最常见的格式：ONNX、TorchScript、SavedModel。

2.1 为什么需要模型导出？

训练框架和推理引擎，本质上是两套东西。训练时我们看重梯度计算、自动微分；推理时只看重前向速度和资源占用。导出就是做这个转换。

  核心目标：
  解耦：把模型从训练框架里解放出来
优化：让推理引擎能做图优化、算子融合、量化
跨平台：一套模型跑在 CPU、GPU、NPU、手机上

2.2 ONNX——通用开放格式

ONNX（Open Neural Network Exchange）是微软和 Meta 联合推出的中间表示。说白了，它就是模型界的「普通话」。不管你是 PyTorch 用户还是 TensorFlow 用户，都能转成 ONNX，然后拿到支持 ONNX 的推理引擎上去跑。

导出流程

以 PyTorch 为例，导出 ONNX 其实就几行代码。但坑很多，我自己就踩过不少。

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

这里有个关键点——dynamic_axes。如果不设动态轴，导出的 ONNX 模型就固定了 batch size 为 1。我在项目中遇到过线上流量波动，不得不重新导出模型的情况，后来才知道加这个参数就行了。

一个小技巧： 导出的 ONNX 模型可以用 onnxruntime 的 Python API 先跑一遍验证，确保精度没掉。

import onnxruntime as ort

session = ort.InferenceSession("resnet18.onnx")
result = session.run(None, {"input": dummy_input.numpy()})
print(result[0].shape)

ONNX 算子兼容性

ONNX 的算子集是有限的。如果你用了 PyTorch 里的某些自定义算子或者太新的算子，导出时可能会报错。嗯，这种情况我遇到过好几次。

常见问题	原因	解决办法
导出报错「Unsupported operator」	算子不在 ONNX 注册表中	用 `torch.onnx.symbolic_opset` 自定义映射
动态 shape 推理失败	模型中某个层不支持动态输入	检查 `dynamic_axes` 设置，或固定输入尺寸
精度微妙偏差	算子实现方式不同	用 `onnxruntime` 对比输出，逐个定位

2.3 TorchScript——PyTorch 自己的中间表示

TorchScript 是 PyTorch 的亲儿子。它可以追踪你的模型计算图，然后序列化成一种独立于 Python 的表示。说白了，就是让你的 PyTorch 模型跑在没有 Python 解释器的环境里。

两种导出方式

TorchScript 有两种生成方式：Tracing 和 Scripting。我刚开始做的时候也搞混过。

Tracing（追踪）：喂一组输入，记录下实际执行的计算路径。速度快，但不支持动态控制流。
Scripting（脚本）：编译你的模型代码，支持 if/else、循环等控制流。更灵活，但需要模型代码是纯 PyTorch 操作。

# Tracing 方式
model = models.resnet18(pretrained=True)
model.eval()
traced_model = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
traced_model.save("resnet18_traced.pt")

# Scripting 方式
scripted_model = torch.jit.script(model)
scripted_model.save("resnet18_scripted.pt")

我记得第一次导出一个带 if 分支的模型，用 Tracing 死活不对，后来查文档才发现该用 Scripting。你想想看，一个 if-else 的分支路径，Tracing 只跑了一条，当然会翻车。

注意： TorchScript 对 Python 标准库和第三方库的依赖非常敏感。模型里如果用了 numpy、pandas 等非 torch 操作，Scripting 很可能失败。

2.4 SavedModel——TensorFlow 的标准格式

SavedModel 是 TensorFlow 推荐的模型导出格式。它把模型的计算图、权重、签名信息都打包在一个目录里。结构清晰，适合 TensorFlow Serving 这样的生产系统。

导出示例

用 TensorFlow 2.x 的 Keras API 导出很简单。

import tensorflow as tf

model = tf.keras.applications.ResNet50(weights="imagenet")
model.build((None, 224, 224, 3))
model.save("resnet50_savedmodel", save_format="tf")

导出的目录结构长这样：

resnet50_savedmodel/
├── assets/
├── variables/
│   ├── variables.data-00000-of-00001
│   └── variables.index
├── saved_model.pb
└── keras_metadata.pb

saved_model.pb 里存的是计算图的 protobuf 表示。variables 目录里是权重文件。我真见过有人只拷贝了 .pb 文件，忘记带 variables，上线时找不到权重——大家千万别犯这个错。

签名定义： SavedModel 支持多签名，你可以同时定义「推理」「训练」「特征提取」等入口。

@tf.function(input_signature=[tf.TensorSpec(shape=[None,224,224,3], dtype=tf.float32)])
def infer(x):
    return model(x, training=False)

tf.saved_model.save(model, "resnet50_with_signature", signatures={"infer": infer})

三种格式的对比

什么时候用哪个？我个人的经验是这样选的：

格式	优点	缺点	适用场景
ONNX	跨框架通用，生态好	算子兼容性有限	需要跨平台部署、多框架切换
TorchScript	和 PyTorch 无缝集成	只能被 C++/Python 的 libtorch 加载	纯 PyTorch 技术栈
SavedModel	生产环境成熟，支持版本管理	对非 TensorFlow 生态支持弱	TensorFlow Serving、TFX 流水线

知识结构总览

下面这张图帮你理清三者的关系和选择逻辑：

避坑指南与经验总结

做了这么多年模型导出，我归纳了三条最实用的小贴士：

永远做精度验证：导出后的模型和原始模型在同一个输入上跑一次，对比输出差。差超过 1e-4 就要排查。我曾经被一个 float16 精度问题坑了三天，从此养成习惯。
注意动态 shape：很多模型导出时默认 batch size=1，生产环境一变就报错。导出时就要把 dynamic_axes 或 input_signature 设好。
版本对齐：ONNX opset、PyTorch 版本、onnxruntime 版本，这三者一定要兼容。我遇到过 opset 太高导致 onnxruntime 不认的情况，最后只能降版本。

一个容易忽略的点： 模型里的预处理逻辑（归一化、resize、通道变换）要不要包含在导出模型里？我的建议是分离出来，让前处理在 Python 或 C++ 层面完成。这样模型更纯净，不同框架之间切换也更方便。

好了，模型导出这部分就先聊到这儿。你想想看，这三种格式各有千秋，没有绝对的「最好」。选哪个，取决于你的技术栈和目标平台。只要能顺利把模型送进推理引擎，就算第一步走对了。