第4章 ONNX模型导出与优化

第4章：ONNX模型导出与优化

说实话，ONNX 这东西刚出来那会儿，我其实挺不以为然的。觉得不就是个中间格式嘛，能有多大用？直到有一次，我在一个嵌入式项目里，需要把 PyTorch 模型部署到 NVIDIA Jetson 上，才发现——没有 ONNX，你根本玩不转。

ONNX 说白了，就是深度学习模型的「通用语言」。你拿 PyTorch 训练好的模型，转成 ONNX，就能在各种推理引擎上跑。TensorRT、OpenVINO、ONNX Runtime，都认它。今天我就带你走一遍完整的流程：从 PyTorch 导出 ONNX，到检查算子兼容性，再到用 Simplifier 优化，最后搞定动态轴和静态轴。

核心流程： PyTorch 模型 → torch.onnx.export() → ONNX 模型 → 算子兼容性检查 → ONNX Simplifier 优化 → 动态/静态轴设置 → 部署就绪

4.1 PyTorch 模型转 ONNX：没那么玄乎

导出 ONNX，PyTorch 官方给了个接口叫 torch.onnx.export()。用法其实挺简单的，但坑不少。我刚开始用的时候，就踩过一个——忘了设 opset_version，结果导出的模型在 TensorRT 上死活跑不起来。

给你看个标准例子：

import torch
import torchvision.models as models

# 加载一个预训练的 ResNet-18
model = models.resnet18(pretrained=True)
model.eval()

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

# 导出 ONNX
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"}
    }
)

嗯，这里有几个关键点我得跟你聊聊。

opset_version 怎么选？ 我个人习惯用 11 或 13。版本太低，有些新算子不支持；版本太高，老推理引擎又不认。你想想看，要是你的目标平台是 Jetson Xavier，它默认的 TensorRT 版本可能只支持到 opset 11。选高了，导出没问题，但推理引擎加载时会报错。

小技巧： 不确定目标平台支持哪个 opset？去 ONNX 官网查算子兼容性矩阵，或者直接试——先导出一个 opset 11 的，跑通了再升级。

do_constant_folding 这个参数，我建议始终设为 True。它会把模型里那些固定不变的运算（比如权重乘常数）提前算好，减少推理时的计算量。说白了就是帮你省点电，在嵌入式设备上，省电就是省命。

4.2 ONNX 算子兼容性检查：别等部署了才发现问题

模型导出成功，不代表万事大吉。我曾经有个项目，模型在 PyTorch 里跑得飞起，转成 ONNX 后，在 ONNX Runtime 上一跑——直接崩了。查了半天，发现是用了 torch.nn.functional.grid_sample，这个算子在 ONNX 里支持得不太好。

所以，导出后一定要做算子兼容性检查。怎么查？两个办法：

用 ONNX Runtime 直接跑一遍——最实在的方法。能跑通，说明算子基本兼容。
用 onnx.checker 做静态检查——检查模型结构是否合法。

import onnx

# 加载模型
model = onnx.load("resnet18.onnx")

# 静态检查
onnx.checker.check_model(model)
print("模型结构检查通过！")

# 用 ONNX Runtime 跑推理验证
import onnxruntime as ort
import numpy as np

ort_session = ort.InferenceSession("resnet18.onnx")
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

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

# 跑推理
outputs = ort_session.run([output_name], {input_name: dummy_input})
print("推理验证通过！输出形状：", outputs[0].shape)

注意： 静态检查通过，不代表推理时没问题。有些算子结构上合法，但运行时可能因为数据形状或数值范围出问题。所以「静态检查 + 动态推理验证」两个都要做。

如果遇到不兼容的算子怎么办？我一般这么处理：

查 ONNX 官方文档，看有没有替代算子
用 torch.onnx.register_custom_op 注册自定义算子
实在不行，改模型结构，换一个 PyTorch 里功能类似但 ONNX 支持更好的算子

4.3 ONNX Simplifier 优化：给模型减减肥

导出的 ONNX 模型，往往有很多冗余结构。比如一些恒等映射、没用的节点、重复的计算。这些在 PyTorch 动态图里无所谓，但在 ONNX 静态图里，就是累赘。

ONNX Simplifier 就是干这个的。它能帮你：

删除恒等节点（Identity ops）
合并连续的 reshape 操作
消除死代码（永远不会执行的节点）
常量折叠（把能提前算的常数运算算掉）

用法很简单：

pip install onnx-simplifier

python -m onnxsim resnet18.onnx resnet18_simplified.onnx

或者用 Python 代码：

import onnx
from onnxsim import simplify

model = onnx.load("resnet18.onnx")
model_simplified, check = simplify(model)

if check:
    onnx.save(model_simplified, "resnet18_simplified.onnx")
    print("简化成功！")
else:
    print("简化失败，请检查模型结构。")

我记得有一次，一个 200MB 的模型，Simplifier 跑完后变成了 180MB。别小看这 20MB，在嵌入式设备上，内存就是钱。而且推理速度也快了大概 15%。

经验之谈： Simplifier 不是万能的。有些模型简化后反而会出问题，比如某些自定义算子被误删。所以简化后一定要重新跑一遍推理验证，确保精度没掉。

4.4 动态轴与静态轴设置：灵活性与性能的博弈

这个点，我觉得是 ONNX 导出里最容易被忽视的。什么叫动态轴？说白了，就是允许输入的形状在推理时变化。比如 batch size 可以是 1、4、8，或者图片尺寸可以是 224x224、256x256。

静态轴呢？就是固定死了。输入必须是 1x3x224x224，多一点少一点都不行。

你可能会想：「那肯定用动态轴啊，灵活！」嗯，但代价是性能。动态轴会让推理引擎做更多的内存分配和形状推导，速度会慢一些。静态轴则可以做更多编译期优化，跑得更快。

怎么选？我个人的原则是：

场景	推荐设置	原因
服务端推理（批量处理）	动态 batch size	需要灵活处理不同请求量
嵌入式设备（固定输入）	全静态	性能优先，内存占用可控
视频流处理（固定分辨率）	静态分辨率 + 动态 batch	分辨率固定，但可能多路并行
多任务模型	按需设置	不同输出可能有不同动态需求

设置动态轴的方法，就是在 torch.onnx.export() 里传 dynamic_axes 参数。刚才的代码里已经演示过了。我再给你拆解一下：

dynamic_axes = {
    "input": {0: "batch_size", 2: "height", 3: "width"},
    "output": {0: "batch_size"}
}

这里的意思是：输入的维度 0（batch）是动态的，维度 2（高）和维度 3（宽）也是动态的。输出的维度 0 跟着 batch 变。

坑点提醒： 动态轴不是你想设就能设的。有些算子（比如全连接层）要求输入形状固定，设了动态轴反而会报错。我遇到过最头疼的一次，是模型里有个 torch.nn.Flatten()，它会把多维输入压平，动态轴一设，压平后的维度就乱了。解决办法是改用 torch.reshape() 手动控制。

最后给你个建议：先静态，再动态。先把模型用静态轴导出，跑通验证。确认没问题了，再逐步放开动态轴。这样出问题了，你能快速定位是动态轴的问题，还是模型本身的问题。

我的习惯： 在项目初期，所有轴都设成静态。等模型稳定了，再根据实际部署需求，只放开必要的动态轴。少即是多。

好了，ONNX 导出与优化这块，核心就是这些。从 PyTorch 到 ONNX，再到 Simplifier 瘦身，最后搞定动态轴——每一步都有坑，但每一步也都有解法。你按这个流程走一遍，基本不会出大问题。

📚 课程目录