4. 模型转换基础：ONNX作为中间格式

模型转换这事儿，说白了就是让不同框架之间能「说上话」。PyTorch 训练好的模型，想部署到 TensorRT 上跑？或者想在 OpenVINO 上推理？中间必须有个通用格式来搭桥。ONNX 就是这个桥。

我个人习惯把 ONNX 看作模型的「通用语言」。不管你是 PyTorch 出身，还是 TensorFlow 老手，最终都得翻译成 ONNX，才能被各种推理引擎理解。今天咱们就聊聊这个翻译过程里，那些容易踩的坑。

4.1 为什么选 ONNX？

ONNX 全称是 Open Neural Network Exchange。微软和 Facebook 搞出来的。它不依赖任何训练框架，只描述计算图的结构和权重。你想想看，一个模型在 PyTorch 里是 nn.Module，在 TensorFlow 里是 SavedModel，到了 ONNX 里就变成了统一的 GraphProto。

我在项目中遇到过最头疼的事：客户给的模型是 PyTorch 的，但推理平台只支持 TensorRT。没有 ONNX 的话，我得手写网络结构再加载权重，那工作量...嗯，想想就头大。

核心优势：

  • 框架无关：一次导出，到处运行
  • 算子标准化：Conv、Relu、BatchNorm 等都有统一定义
  • 支持量化：INT8、FP16 都能在 ONNX 层面做
  • 生态完善：TensorRT、OpenVINO、ONNX Runtime 都支持

4.2 从 PyTorch 导出 ONNX

PyTorch 导出 ONNX 用的是 torch.onnx.export。这个函数看起来简单，但参数调起来能让你怀疑人生。

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",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有几个关键点我得提醒你：

  • model.eval() 必须加：不切到 eval 模式，BatchNorm 和 Dropout 的行为会乱掉。我曾经因为这个排查了一整天，最后发现是忘了加这一行。
  • dummy_input 的 shape 要对：ONNX 导出时会做一次推理来 trace 计算图。输入 shape 不对，trace 出来的图就是错的。
  • opset_version 选 11 或更高：低版本 opset 不支持很多新算子。我一般用 11，兼容性和功能都比较平衡。
  • dynamic_axes 要谨慎：如果你需要动态 batch size，就加上。但有些推理引擎对动态 shape 支持不好，反而会出问题。

小技巧：导出后一定要用 onnx.checker.check_model 验证一下。我每次导出完都会跑一遍这个检查，能发现很多隐藏问题。

4.3 从 TensorFlow 导出 ONNX

TensorFlow 导出 ONNX 稍微麻烦点。官方推荐用 tf2onnx 这个工具。有两种方式：

方式一：命令行转换

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --opset 11

方式二：Python API 转换

import tf2onnx
import tensorflow as tf

model = tf.saved_model.load('./saved_model')
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

output_path = "./model.onnx"
model_proto, _ = tf2onnx.convert.from_function(
    model.signatures['serving_default'],
    input_signature=spec,
    opset=11,
    output_path=output_path
)

我个人更推荐用 Python API，因为可以更精细地控制输入输出的名字和类型。命令行虽然方便，但遇到问题调试起来很痛苦。

注意：TensorFlow 的 NHWC 格式和 ONNX 的 NCHW 格式经常搞混。转换时记得检查一下数据排布。我遇到过好几次，模型在 TF 里跑得好好的，转成 ONNX 后结果全错，最后发现是通道顺序的问题。

4.4 常见转换问题与调试

模型转换从来不是一帆风顺的。我总结了几个高频问题，你大概率会遇到：

4.4.1 算子不支持

这是最常见的问题。PyTorch 或 TensorFlow 里的一些自定义算子，ONNX 里没有对应的实现。比如 torch.nn.functional.grid_sample，在低版本 opset 里就不支持。

解决办法：

  • 升级 opset 版本（试试 13 或更高）
  • 用 ONNX 支持的算子重新实现
  • 写自定义算子（Custom Op），但工作量比较大

4.4.2 动态 shape 问题

导出时指定了动态 batch size，但推理引擎报错说 shape 不匹配。这通常是因为模型里有 reshape 或 view 操作，它们对动态 shape 很敏感。

我曾经...遇到过一个问题：模型里有个 view(-1, 64)，导出时 batch size 是 1，一切正常。但推理时 batch size 变成 4，view 的维度计算就错了。后来改成 reshape 并显式指定维度才解决。

4.4.3 精度差异

转换后模型输出和原始框架有微小差异。这通常是因为：

  • 算子实现细节不同（比如某些激活函数的边界处理）
  • 权重精度变化（FP32 转 FP16 时）
  • 计算顺序不同导致浮点误差累积

调试方法：

  1. 用 onnxruntime 跑一遍 ONNX 模型，和原始框架对比输出
  2. 逐层对比中间结果，定位差异层
  3. 检查该层的算子实现是否有差异
  4. 必要时用 onnx-simplifier 简化模型结构

4.4.4 控制流问题

如果你的模型里有 if-else 或循环（比如 RNN、LSTM），导出 ONNX 时会比较麻烦。ONNX 支持 If 和 Loop 算子，但很多推理引擎不支持。

我建议：尽量把控制流逻辑放到模型外部。比如在预处理或后处理里做判断，模型只做纯计算。这样转换起来省心很多。

4.5 调试工具推荐

工具 用途 我的评价
Netron 可视化 ONNX 模型结构 必备工具，一眼看出网络结构对不对
onnxruntime 本地推理验证 转换后第一件事就是用这个跑一遍
onnx-simplifier 简化模型，去除冗余节点 能解决很多莫名其妙的转换问题
onnx-graphsurgeon 手动修改计算图 高级玩家用的，可以精确控制每个节点

说实话，模型转换是个熟能生巧的活。刚开始你可能觉得处处是坑，但多转几次就摸到门道了。记住一个原则：先简化，再转换。把模型里的花哨操作都去掉，只保留核心计算，转换成功率会高很多。

我的工作流：

导出 ONNX → Netron 查看结构 → onnxruntime 验证精度 → onnx-simplifier 简化 → 再次验证 → 部署到目标平台。每一步都别跳过，省事就是费事。

嗯，关于模型转换的基础就聊这么多。下一章咱们会深入 ONNX 的算子层面，看看怎么手动优化计算图。到时候你会觉得，今天这些坑都只是开胃菜。

← 上一章 📖 返回目录 下一章 →