3、ONNX模型基础与优化：从导出到推理的全流程实战

ONNX，说白了就是深度学习模型的"通用语言"。不管你是用PyTorch、TensorFlow还是其他框架训练的模型，都能转成ONNX格式。我个人习惯把ONNX看作模型部署的中间桥梁——它不负责训练，只负责让模型在不同推理引擎之间自由迁移。

这一章，我会带你走完ONNX的完整流程：从协议标准到推理演示，再到模型简化和节点编辑。嗯，这些都是我在实际项目中反复踩过的坑，今天一并分享给你。

3.1 ONNX协议与算子标准

ONNX（Open Neural Network Exchange）的核心思想很简单：定义一个统一的计算图表示。每个节点代表一个算子（Operator），边代表张量数据流。你想想看，如果每个框架都用自己的格式，那部署工程师就得为每个框架写一套推理代码，多累啊。

ONNX的算子标准分为三个层级：

层级 说明 示例
核心算子 所有框架通用的基础操作 Conv、Relu、MaxPool、Add
扩展算子 特定领域需要的操作 LSTM、GRU、TopK
自定义算子 用户自己实现的特殊操作 YOLO的Decode、NMS

我在项目中遇到过一个问题：YOLOv5的Focus层在ONNX里没有直接对应的算子。当时我查了半天文档，最后发现可以用切片+拼接的组合来实现。这就是算子标准的重要性——你得知道哪些操作是"原生支持"的，哪些需要手动拆解。

避坑指南： 我曾经因为算子版本不匹配，在TensorRT里跑ONNX模型时直接报错。记住，ONNX opset version一定要和推理引擎支持的版本对齐。比如TensorRT 8.x推荐使用opset 11-13，太新或太旧都可能出问题。

3.2 ONNX Runtime推理演示

ONNX Runtime是微软推出的跨平台推理引擎。它支持CPU、GPU、甚至NPU。我个人习惯在导出ONNX后，先用ONNX Runtime跑一遍，确认模型没问题再交给TensorRT。

下面是一个完整的推理代码示例：

import onnxruntime as ort
import numpy as np
from PIL import Image

# 1. 创建推理会话
session = ort.InferenceSession("yolov5s.onnx")

# 2. 获取输入输出信息
input_name = session.get_inputs()[0].name
output_names = [output.name for output in session.get_outputs()]
print(f"输入: {input_name}, 形状: {session.get_inputs()[0].shape}")
print(f"输出: {output_names}")

# 3. 预处理图像
image = Image.open("test.jpg").resize((640, 640))
input_tensor = np.array(image).astype(np.float32) / 255.0
input_tensor = np.transpose(input_tensor, (2, 0, 1))  # HWC -> CHW
input_tensor = np.expand_dims(input_tensor, axis=0)   # 添加batch维度

# 4. 执行推理
outputs = session.run(output_names, {input_name: input_tensor})

# 5. 后处理
print(f"推理结果形状: {[o.shape for o in outputs]}")

你可能会问：为什么要先用ONNX Runtime？原因有三：

  • 验证模型正确性——确保导出过程没出问题
  • 对比精度——ONNX Runtime的结果可以作为基准
  • 调试方便——ONNX Runtime的错误信息比TensorRT友好得多
小技巧： 我建议在ONNX Runtime里设置providers=['CUDAExecutionProvider', 'CPUExecutionProvider']，这样GPU可用时自动用GPU，否则回退到CPU。生产环境里这个配置很实用。

3.3 onnx-simplifier模型简化

导出的ONNX模型往往包含很多冗余操作。比如PyTorch导出的模型里，经常出现一堆Shape、Gather、Unsqueeze节点。这些节点在推理时完全没用，只会拖慢速度。

onnx-simplifier就是干这个的——它会把计算图中那些"只参与形状计算"的节点全部干掉。

使用方法超级简单：

# 安装
pip install onnx-simplifier

# 命令行简化
python -m onnxsim yolov5s.onnx yolov5s_simplified.onnx

# Python API
import onnx
from onnxsim import simplify

model = onnx.load("yolov5s.onnx")
model_simplified, check = simplify(model)
onnx.save(model_simplified, "yolov5s_simplified.onnx")
print(f"简化成功: {check}")

我记得第一次用onnx-simplifier时，一个原本有1200个节点的YOLOv5模型，简化后只剩400多个节点。推理速度提升了将近30%。你想想看，这优化效果多明显。

核心要点： onnx-simplifier主要做三件事：常量折叠、形状推理、死节点删除。它不会改变模型的数学结果，只去掉冗余计算。每次导出ONNX后，我第一件事就是跑simplifier。

3.4 onnx-graphsurgeon节点编辑

onnx-simplifier能解决大部分问题，但有些场景需要更精细的控制。比如你想把YOLO的Decode层合并成一个自定义节点，或者你想修改某个算子的属性。这时候就需要onnx-graphsurgeon出场了。

onnx-graphsurgeon是NVIDIA推出的ONNX图编辑工具。它允许你直接操作计算图的节点、张量和属性。

来看一个实际例子：把Conv+BN+Relu三个节点合并成一个节点：

import onnx_graphsurgeon as gs
import onnx

# 加载模型
graph = gs.import_onnx(onnx.load("model.onnx"))

# 找到Conv+BN+Relu的序列
for node in graph.nodes:
    if node.op == "Conv":
        bn_node = node.o()  # 获取输出节点
        if bn_node and bn_node.op == "BatchNormalization":
            relu_node = bn_node.o()
            if relu_node and relu_node.op == "Relu":
                # 创建融合节点
                fused_node = gs.Node(
                    op="ConvBnRelu",
                    inputs=node.inputs,
                    outputs=relu_node.outputs,
                    attrs={**node.attrs, **bn_node.attrs}
                )
                # 替换原节点
                graph.nodes.remove(node)
                graph.nodes.remove(bn_node)
                graph.nodes.remove(relu_node)
                graph.nodes.append(fused_node)

# 清理并导出
graph.cleanup()
onnx.save(gs.export_onnx(graph), "model_fused.onnx")

我在项目中遇到过更复杂的情况：YOLOv5的Detect层包含多个输出分支，我需要把它们合并成一个输出，方便TensorRT处理。用onnx-graphsurgeon，我花了大概2小时就搞定了。如果用其他方法，可能得折腾一整天。

注意： onnx-graphsurgeon操作的是计算图结构，一旦改错，模型可能直接报废。我建议每次修改前都备份原始模型，并且修改后一定要用ONNX Runtime验证结果一致性。

3.5 动态batch与静态batch转换

这是部署中最容易踩坑的地方。静态batch意味着模型输入形状固定，比如[1, 3, 640, 640]。动态batch则允许输入形状变化，比如[?, 3, 640, 640]。

两种方式的对比：

特性 静态batch 动态batch
推理速度 快（引擎优化充分） 稍慢（需要动态分配）
内存占用 固定 可变
灵活性 低 高
TensorRT支持 完美 有限制

在PyTorch中导出动态batch的ONNX模型：

import torch

# 定义动态轴
dynamic_axes = {
    "input": {0: "batch_size"},  # batch维度动态
    "output": {0: "batch_size"}
}

# 导出
torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=dynamic_axes,
    opset_version=13
)

我个人习惯的做法是：开发阶段用动态batch，方便调试；生产部署时转成静态batch，追求极致性能。你想想看，线上服务的请求量是固定的，没必要为动态batch付出性能代价。

经验之谈： 如果你一定要用动态batch，记得在TensorRT里设置opt_batch_size和max_batch_size。我曾经因为没设置这两个参数，导致模型在batch=16时直接OOM。嗯，那次教训挺深刻的。

本章知识体系

下面这张图展示了ONNX模型处理的核心流程：

ONNX模型处理流程 ONNX模型处理核心流程 原始模型 PyTorch/TF ONNX导出 torch.onnx.export 模型简化 onnx-simplifier 节点编辑 graphsurgeon 动态batch 灵活但性能略低 静态batch 性能最优 ONNX Runtime验证 精度对齐检查 TensorRT部署 最终推理引擎 💡 核心原则：先验证，再优化 每一步都要用ONNX Runtime做精度对齐

从这张图你能看到，ONNX模型处理不是一步到位的。它需要经过导出、简化、编辑、batch策略选择，最后才能交给TensorRT。每一步都有对应的工具和注意事项。

嗯，这一章的内容就到这里。记住一句话：ONNX是桥梁，不是终点。真正考验部署工程师功力的，是如何把这座桥搭得又快又稳。

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