第3章模型格式基础

3、模型格式基础：ONNX格式详解、PyTorch模型导出为ONNX、ONNX算子兼容性检查

做TensorRT-LLM部署，绕不开ONNX这个中间格式。说白了，ONNX就是模型界的“通用语言”。你训练好的PyTorch模型，TensorRT-LLM不认识，得先翻译成ONNX，它才能理解。

我个人习惯把ONNX比作“翻译官”。它夹在PyTorch和TensorRT中间，负责把PyTorch那一套算子，转成TensorRT能吃的格式。这个环节要是出问题，后面全白搭。

3.1 ONNX格式到底长什么样？

ONNX全称Open Neural Network Exchange，微软和Facebook联合搞的。它的核心是一个计算图，包含三要素：节点（算子）、边（张量）、图属性。

我刚开始接触ONNX时，总觉得它很神秘。后来用Netron工具打开一个ONNX文件，发现其实就是个有向无环图。每个节点代表一个操作，比如Conv、Relu、Add。边就是数据流动的通道。

ONNX文件本身是protobuf序列化的二进制文件。你可以用Python读取它：

import onnx

model = onnx.load("model.onnx")
print(onnx.helper.printable_graph(model.graph))

这段代码会打印出整个计算图的结构。嗯，这里要注意，ONNX的算子版本很重要。不同版本的算子，支持的属性可能不一样。

ONNX的核心数据结构：

ModelProto：最外层，包含模型版本、生产者信息、图结构
GraphProto：计算图本体，包含节点、初始值、输入输出
NodeProto：单个算子节点，包含算子类型、输入输出名称、属性
TensorProto：张量数据，用于存储权重和偏置

我建议你养成一个好习惯：每次导出ONNX后，先用Netron看一眼图结构。我在项目中遇到过，明明模型跑得好好的，导出ONNX后多了一个莫名其妙的Reshape节点，导致TensorRT推理失败。

3.2 PyTorch模型导出为ONNX

PyTorch导出ONNX，主要靠torch.onnx.export这个函数。看起来简单，但坑不少。

先看一个标准流程：

import torch
import torch.onnx

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(3, 64, 3)
        self.relu = torch.nn.ReLU()
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

model = MyModel()
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=17
)

这里有几个关键点，我踩过坑，你记一下：

导出ONNX的避坑指南：

model.eval()必须调用：不调用的话，BatchNorm和Dropout的行为会不一样。我曾经因为这个，导出的ONNX在推理时结果对不上。
dummy_input的shape要准确：它决定了ONNX图的输入形状。如果后面要支持动态batch，记得配dynamic_axes。
opset_version别太低：我建议用17以上。太低的版本，很多新算子不支持，或者行为有差异。
input_names/output_names要起好：TensorRT-LLM绑定输入输出时，靠的就是这个名字。起个有意义的，别用默认的"input.1"。

为什么会这样？因为PyTorch的JIT tracer在追踪计算图时，只能看到实际执行的分支。如果你的模型里有if-else，没走到的分支就不会出现在ONNX里。所以，导出时一定要确保dummy_input能覆盖所有路径。

3.3 ONNX算子兼容性检查

模型导出成ONNX，不代表TensorRT-LLM就能吃。算子兼容性，是部署中最头疼的问题之一。

我遇到过最典型的情况：模型里用了torch.nn.functional.grid_sample，PyTorch导出ONNX没问题，但TensorRT不支持这个算子。结果就是，推理时直接报错。

怎么检查兼容性？我一般用三步：

用onnxruntime跑一遍：先确认ONNX本身没问题。如果onnxruntime都跑不通，那肯定是导出环节出错了。
用polygraphy检查算子：NVIDIA官方工具，能列出ONNX中所有算子，并标注哪些TensorRT不支持。
手动查TensorRT算子文档：虽然麻烦，但最靠谱。特别是自定义算子，一定要确认。

polygraphy的使用方法很简单：

polygraphy inspect model model.onnx --mode=basic

这条命令会输出ONNX模型的所有算子列表。你对照TensorRT的算子支持表，一个一个看。

常见的不兼容算子：

算子	问题	解决方案
torch.nn.functional.grid_sample	TensorRT不支持	用插件实现，或者改模型结构
torch.nn.functional.interpolate	部分模式不支持	改用resize算子，或指定支持的插值模式
torch.where	动态shape下可能出问题	确保输入shape固定，或用mask替代
torch.sort / torch.topk	TensorRT支持有限	尽量用固定k值，避免动态排序

你想想看，如果模型里有这些算子，导出ONNX时不会报错，但到TensorRT推理时就卡住了。排查起来特别费时间。

我个人习惯，在模型设计阶段就考虑算子兼容性。能用Conv+Relu解决的问题，就别用花里胡哨的自定义算子。部署时少一个坑，就多一分安稳。

3.4 本章知识体系

下面这张图，是我总结的ONNX格式适配的核心流程。从PyTorch模型到TensorRT-LLM，中间要经过ONNX导出、算子检查、格式转换三步。

这张图把整个流程串起来了。你从PyTorch模型出发，经过ONNX导出，然后做算子兼容性检查。如果发现不兼容的算子，要么改模型结构，要么写TensorRT插件。检查通过后，才能进入TensorRT-LLM的格式转换环节。

嗯，这里要特别强调：算子检查这一步，千万别跳过。我见过太多人，ONNX导出来就直接扔给TensorRT，结果跑不起来，回头排查半天才发现是算子问题。

我的个人经验总结：

导出ONNX时，opset_version用17以上，避免算子版本问题
每次导出后，用Netron看一眼图结构，确认没有多余节点
算子检查用polygraphy，比手动翻文档快得多
遇到不兼容算子，优先考虑改模型结构，写插件是最后的选择

ONNX格式这块，说白了就是“翻译”工作。翻译得好，后面一路顺畅。翻译得不好，后面全是坑。你花30分钟把ONNX导出和检查做好，能省下后面3小时的排查时间。

我曾经在一个项目里，因为ONNX导出时忘了设dynamic_axes，导致模型只能跑固定batch size。后来线上流量波动，不得不重新导出、重新转换，折腾了一整天。从那以后，我每次导出ONNX都会检查dynamic_axes配置。

好了，ONNX格式的基础就讲到这里。记住：导出、检查、转换，三步缺一不可。