3、模型序列化基础：ONNX格式介绍、PyTorch模型导出为ONNX、TensorFlow模型导出为ONNX

好，咱们进入第三章。说实话，模型序列化这块，是很多同学在落地部署时踩的第一个坑。你模型训得再好，导不出来，或者导出来精度不对，那前面全白干。

ONNX，全称是 Open Neural Network Exchange。说白了，它就是深度学习模型的“通用语言”。你想想看，PyTorch 有自己的一套，TensorFlow 也有自己的一套，各家框架各说各话。ONNX 就是那个翻译官，让不同框架的模型能互相串门。

我个人习惯，把 ONNX 看作是模型部署的“中间态”。训练用 PyTorch，部署用 TensorRT，中间这个桥梁就是 ONNX。它不参与训练，也不直接做推理加速，但它决定了你的模型能不能顺利跑到目标硬件上。

3.1 ONNX 格式的核心概念

ONNX 文件里到底装了什么？我拆开给你看。

• 计算图（Graph）：描述了模型的所有计算节点和数据流向。每个节点就是一个算子，比如 Conv、Relu、MatMul。
• 权重参数（Initializer）：模型训练好的权重和偏置，以张量形式存储。
• 输入输出（Input/Output）：定义了模型的输入张量名称、数据类型、形状，以及输出张量的信息。
• 元数据（Metadata）：比如模型版本、作者、框架来源等辅助信息。

嗯，这里要注意：ONNX 的算子版本很重要。不同版本的 ONNX 支持的算子不一样。我在项目中遇到过，导出的 ONNX 模型在旧版本 TensorRT 上跑不了，就是因为用了高版本才有的算子。

3.2 PyTorch 模型导出为 ONNX

PyTorch 导出 ONNX 用的是 torch.onnx.export 这个接口。看起来简单，但里面的门道不少。

先看一个标准流程：

import torch
import torchvision.models as models

# 1. 加载模型，切换到 eval 模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出 ONNX
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    "resnet18.onnx",          # 输出文件名
    export_params=True,       # 导出权重参数
    opset_version=11,         # ONNX 算子版本
    do_constant_folding=True, # 常量折叠优化
    input_names=['input'],    # 输入名称
    output_names=['output'],  # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

print("ONNX 导出完成！")

这里我重点说几个坑：

• eval 模式必须开：我见过有人忘了 model.eval()，结果导出的 ONNX 里包含了 Dropout 和 BatchNorm 的训练行为，推理结果完全不对。
• dummy_input 的形状要对：它决定了 ONNX 模型的输入形状。如果你后续要支持动态 batch，记得配 dynamic_axes。
• do_constant_folding：这个建议打开。它会把一些固定计算提前算好，减少推理时的计算量。

导出后怎么验证？我一般用 ONNX Runtime 跑一遍，对比 PyTorch 的输出：

import onnx
import onnxruntime as ort
import numpy as np

# 检查 ONNX 模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)

# 用 ONNX Runtime 推理
ort_session = ort.InferenceSession("resnet18.onnx")
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# 构造输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
ort_outputs = ort_session.run([output_name], {input_name: input_data})

# 对比 PyTorch 输出
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data))

# 检查误差
diff = np.abs(ort_outputs[0] - torch_outputs.numpy()).max()
print(f"最大误差：{diff:.6f}")

如果误差在 1e-5 以内，基本没问题。超过 1e-3 就要警惕了，可能是算子精度问题。

3.3 TensorFlow 模型导出为 ONNX

TensorFlow 导出 ONNX 稍微麻烦一点。因为 TensorFlow 的图结构更复杂，而且有 SavedModel、Keras H5、Checkpoint 等多种格式。

我推荐用 tf2onnx 这个工具。它专门干这个事，而且社区活跃，踩过的坑都填得差不多了。

先安装：

pip install tf2onnx

然后看代码：

import tensorflow as tf
import tf2onnx

# 1. 加载 Keras 模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')
model.summary()

# 2. 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 3. 导出 ONNX
output_path = "mobilenetv2.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=13,
    output_path=output_path
)

print(f"ONNX 模型已保存到 {output_path}")

如果是 SavedModel 格式，用这个：

import tf2onnx

# 加载 SavedModel
model = tf.saved_model.load("path/to/saved_model")

# 获取推理函数
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 转换
tf2onnx.convert.from_function(
    model.signatures['serving_default'],
    input_signature=spec,
    output_path="model.onnx"
)

这里有个关键点：TensorFlow 的 NHWC 格式和 ONNX 的 NCHW 格式不一样。你想想看，ONNX 默认用 NCHW，但 TensorFlow 习惯用 NHWC。如果不做处理，导出的 ONNX 模型在推理时可能会出问题。

验证方法和 PyTorch 类似，用 ONNX Runtime 跑一遍，对比 TensorFlow 的输出：

import onnxruntime as ort
import numpy as np

ort_session = ort.InferenceSession("mobilenetv2.onnx")
input_name = ort_session.get_inputs()[0].name

# 构造输入（注意是 NCHW 格式）
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# ONNX Runtime 推理
ort_outputs = ort_session.run(None, {input_name: input_data})

# TensorFlow 推理（需要转回 NHWC）
tf_input = np.transpose(input_data, (0, 2, 3, 1))
tf_output = model.predict(tf_input)

# 对比
diff = np.abs(ort_outputs[0] - tf_output).max()
print(f"最大误差：{diff:.6f}")

3.4 导出后的常见问题与排查

导出 ONNX 不是终点，而是起点。我整理了几个高频问题：

问题现象	可能原因	解决方案
ONNX 模型无法加载	算子版本不兼容	降低 opset_version，或升级 ONNX 库
推理结果偏差大	模型未切换到 eval 模式	导出前加 model.eval()
动态 batch 不生效	dynamic_axes 配置错误	检查输入输出名称是否匹配
TensorFlow 导出失败	使用了自定义算子	用 tf2onnx 的 custom_ops 参数注册
ONNX 文件过大	包含了不必要的节点	用 onnx-simplifier 简化模型

下一章，我们会深入 ONNX 模型的优化技巧，包括如何剪枝、如何合并算子、如何消除冗余节点。这些都是为 TensorRT 加速铺路的硬功夫。