3. ONNX模型基础

核心原理：ONNX作为模型交换的“通用语言”

ONNX（Open Neural Network Exchange）本质上是一种中间表示格式，它不直接执行推理，而是定义了一套跨框架的计算图描述规范。当我们将PyTorch模型导出为ONNX时，实际上是在做一次“翻译”：把PyTorch的动态图执行逻辑，转换为一张静态的计算图，图中每个节点代表一个算子（如Conv、Relu、Gemm），每条边代表张量（Tensor）的流动方向。

理解ONNX的关键在于区分三个层次：模型级（ModelProto）、图级（GraphProto）和节点级（NodeProto）。一个ONNX文件就是一个序列化的ModelProto，它包含一个GraphProto；GraphProto由多个NodeProto组成节点列表，同时维护了图的输入输出张量信息。这种嵌套结构使得ONNX既能表达简单的线性网络，也能表达带有分支、循环的复杂拓扑。

核心要点：ONNX文件不是“模型权重包”，而是“计算图+权重”的复合体。权重以初始化器（Initializer）的形式存放在图中，作为节点的常量输入。因此，ONNX文件的大小主要由权重张量决定，而计算图本身通常只占很小体积。

模型结构：从ProtoBuf到可视化计算图

ONNX基于Google Protobuf序列化，因此文件后缀为.onnx。我们可以用Python API快速解析其结构：

import onnx
model = onnx.load("resnet18.onnx")
graph = model.graph
print("节点数量:", len(graph.node))
print("输入:", [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in graph.input])
print("输出:", [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim]) for o in graph.output])
print("初始化器数量:", len(graph.initializer))

一个典型的ONNX图包含以下核心字段：

字段作用示例
node算子节点列表，按拓扑序排列Conv、Relu、MaxPool
input图的输入张量定义（含形状与类型）input: float32[1,3,224,224]
output图的输出张量定义output: float32[1,1000]
initializer权重与常量张量（不参与梯度）conv1.weight: float32[64,3,7,7]
value_info中间张量的形状推断结果layer1_out: float32[1,64,56,56]

每个节点本身又包含op_type（算子类型）、input（输入张量名列表）、output（输出张量名列表）以及attribute（算子属性，如卷积的kernel_shape、strides）。张量在图中通过名称字符串进行关联，而非指针或索引，这保证了图的序列化友好性。

节点与张量：理解数据流的关键

ONNX中的张量（Tensor）是不可变的数据载体，它只有名称、数据类型和形状三个核心属性。节点则是计算单元，它消费输入张量并产生输出张量。这种设计使得ONNX图天然适合静态分析——我们可以不执行计算，仅通过遍历节点和张量名就能还原完整的数据流路径。

以ResNet中的一个BasicBlock为例，其ONNX子图结构如下：

节点序列（简化）:
  Conv_0: input[input, res1.weight] → output[conv1_out]
  Relu_1: input[conv1_out] → output[relu1_out]
  Conv_2: input[relu1_out, res2.weight] → output[conv2_out]
  Add_3:  input[conv2_out, input] → output[add_out]  # 残差连接
  Relu_4: input[add_out] → output[block_out]

注意Add_3节点直接引用了图的原始输入input，这就是残差连接在ONNX中的表达方式——通过张量名实现跨层引用。理解这一点对后续TensorRT优化至关重要，因为TensorRT的层融合（如将Conv+BN+Relu融合为单个算子）正是基于这种数据流依赖关系进行重写的。

实操建议：使用onnx.graphsurgeon或netron可视化工具查看模型时，重点关注张量名的命名规律。PyTorch导出的模型张量名通常为"/layer1.0/conv1/Conv_output_0"这种带路径前缀的格式，便于溯源。

动态轴：从固定形状到可变批量

ONNX默认使用静态形状描述张量，即每个维度都是固定整数。但实际部署中，我们往往需要支持不同的batch size或输入分辨率。为此，ONNX引入了动态轴（Dynamic Axes）机制，在导出时通过dynamic_axes参数指定哪些维度是可变的：

import torch
# 导出时允许batch维度和宽高维度动态变化
torch.onnx.export(
    model, dummy_input, "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    },
    opset_version=17
)

导出后，ONNX图中对应维度的dim_param字段会变为字符串（如"batch_size"），而非整数。TensorRT在解析这类模型时，会将其转换为优化配置文件（Optimization Profile）中的动态范围。需要注意的是，并非所有算子都支持动态形状——例如某些全局池化或reshape操作可能要求输入形状完全确定。

常见误区：动态轴并非越多越好。每个动态维度都会增加TensorRT的优化搜索空间，导致构建时间变长。如果业务场景中batch size固定为1，建议导出时保持静态，仅在需要弹性伸缩时才启用动态轴。

案例：将PyTorch模型导出为ONNX并检查动态轴

下面演示一个完整的导出与验证流程。假设我们有一个简单的CNN分类模型：

import torch
import torch.nn as nn
import onnx

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.fc = nn.Linear(16*8*8, 10)
    def forward(self, x):
        x = torch.relu(self.conv(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = SimpleCNN().eval()
dummy = torch.randn(1, 3, 8, 8)

# 导出：batch维度动态，图像尺寸固定
torch.onnx.export(model, dummy, "simple_cnn.onnx",
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
                  opset_version=17)

# 验证动态轴
onnx_model = onnx.load("simple_cnn.onnx")
for inp in onnx_model.graph.input:
    for dim in inp.type.tensor_type.shape.dim:
        print(f"输入维度: dim_value={dim.dim_value}, dim_param={dim.dim_param}")

输出结果中，batch维度的dim_param为"batch"，而其他维度保持dim_value为固定值。这表示TensorRT在转换时，需要为batch维度设置最小/最优/最大三个值（如1/8/32），而图像尺寸则固定为8×8。

误区与复盘：ONNX模型优化的常见陷阱

  • 误区一：忽略opset版本。不同opset版本支持的算子集合不同，TensorRT对高版本opset的支持可能滞后。建议使用opset 13~17之间的版本，兼顾算子覆盖率和兼容性。
  • 误区二：将动态轴用于所有维度。例如将通道数也设为动态，这会导致TensorRT无法进行通道维度的内存布局优化，性能大幅下降。
  • 误区三：不检查中间张量的精度。ONNX默认使用float32，但某些框架导出时可能混入float64算子，导致TensorRT拒绝转换。导出后应使用onnx.checker.check_model和onnx.shape_inference.infer_shapes进行完整性验证。
  • 误区四：忽视initializer的存储布局。ONNX中权重默认按行主序（row-major）存储，而TensorRT内部可能使用列主序或特定向量化布局。转换器会自动处理重排，但若手动修改ONNX权重，则需格外小心。

复盘总结：一个健壮的ONNX导出流程应包含以下步骤：① 使用torch.onnx.export并显式指定opset；② 用onnx.checker验证模型合法性；③ 用onnx.shape_inference推断中间张量形状；④ 用netron可视化检查数据流；⑤ 在TensorRT转换前，先用onnxruntime跑一遍推理，确认输出与PyTorch一致。这五步能过滤掉绝大多数导出问题。

ONNX模型结构解析与TensorRT转换链路 ONNX 文件结构 ModelProto GraphProto node[] initializer input/output 节点示例: Conv(input, W) → conv1 Relu(conv1) → act1 Add(act1, input) → res1 TensorRT 转换链路 ONNX Parser Network Definition Builder + Engine 动态轴处理 (Optimization Profile) 输入张量: input[batch, 3, H, W] batch: min=1, opt=8, max=32 H/W: 固定为 224×224 (静态) 张量数据流与动态轴示意 input [batch,3,224,224] Conv [batch,64,112,112] Relu [batch,64,112,112] MaxPool [batch,64,56,56] Flatten+FC [batch,10] 动态轴传播路径 batch维度: 从input到所有中间张量均保持动态 → TensorRT需为每层分配动态内存池 H/W维度: 经过Conv/Relu保持不变，经过MaxPool减半 → 静态形状可提前优化内存布局 建议: 仅batch维度动态，空间维度固定，以获得最佳推理性能

小结与自查清单

本章我们梳理了ONNX格式的四个核心概念：ProtoBuf序列化结构、图与节点的层次关系、张量名称驱动的数据流以及动态轴的声明与传播规则。在进入TensorRT优化之前，请确保你能够回答以下问题：

  • ONNX文件中的initializer和node.input有什么区别？权重存储在哪一层？
  • 如果两个节点都引用了名为"res2_out"的张量，这代表什么含义？
  • 动态轴在ONNX图中是如何标记的？TensorRT的Optimization Profile需要哪些参数？
  • 为什么说“动态轴越多，TensorRT构建越慢”？

掌握这些基础后，下一阶段我们将聚焦于TensorRT的层融合策略与精度校准方法，届时你会更深刻地理解ONNX计算图结构如何直接影响底层优化决策。

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