4、ONNX解析与模型导入：ONNX格式简介、使用ONNX Parser导入模型、处理动态输入形状、调试ONNX导入失败问题。

好，咱们进入第四章。这一章聊的是ONNX解析和模型导入。说实话，这是整个TensorRT插件开发里最磨人的环节之一。我见过太多人模型写好了，一跑ONNX解析就崩，然后就开始怀疑人生。别急，咱们一步步来。

4.1 ONNX格式简介：它到底是个啥？

ONNX，全称Open Neural Network Exchange。说白了，它就是深度学习模型的“通用语言”。你想想看，PyTorch训练出来的模型，TensorFlow训练出来的模型，它们各自有各自的格式。ONNX就是那个中间人，让不同框架的模型能互相串门。

我个人习惯把ONNX理解成一个计算图。它里面包含了：

  • 节点（Node）：每个节点代表一个算子，比如Conv、Relu、Add。
  • 张量（Tensor）：节点之间的数据流动，就是张量。
  • 图（Graph）：所有节点和张量组成的网络结构。

举个例子，一个简单的卷积网络，ONNX文件里大概长这样：

graph(%input : Float(1, 3, 224, 224)) {
    %conv1_weight : Float(64, 3, 3, 3)
    %conv1_bias : Float(64)
    %conv1_out = Conv(%input, %conv1_weight, %conv1_bias)
    %relu1_out = Relu(%conv1_out)
    %pool1_out = MaxPool(%relu1_out)
    return %pool1_out
}

嗯，这里要注意。ONNX的算子集是标准化的，但不同框架导出的ONNX可能会有细微差别。我在项目中遇到过，PyTorch导出的ONNX里有些算子名字跟TensorRT的不完全一样，这就得靠后面的插件来兜底了。

4.2 使用ONNX Parser导入模型

TensorRT里有个好东西叫ONNX Parser。它的任务就是把ONNX文件“翻译”成TensorRT能理解的网络结构。用法其实挺简单的：

// 创建builder和network
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

// 创建ONNX parser
nvinfer1::IONNXParser* parser = nvonnxparser::createParser(*network, gLogger);

// 解析ONNX文件
if (!parser->parseFromFile("model.onnx", static_cast<int>(ILogger::Severity::kWARNING))) {
    // 解析失败，获取错误信息
    std::cout << "Failed to parse ONNX file!" << std::endl;
    return -1;
}

你看，核心就三步：建网络、建解析器、解析文件。但这里有个坑——parseFromFile返回false不代表所有算子都失败了。它可能只是某个算子没找到，但其他部分已经解析成功了。我刚开始做的时候，一看到返回false就直接放弃，后来才发现可以继续用已经解析好的部分。

小技巧：解析ONNX时，建议把日志级别设成kWARNING或更低。这样你能看到哪些算子被跳过了，哪些被成功转换了。我一般用kINFO，信息更全。

4.3 处理动态输入形状

动态输入形状，说白了就是你的模型输入尺寸不固定。比如你训练时用224x224，但部署时可能遇到192x192或256x256的图片。TensorRT默认是静态的，你得告诉它：“嘿，我这输入是可变的。”

怎么做呢？在ONNX导出时就要设置动态轴：

# PyTorch导出ONNX时设置动态轴
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size', 2: 'height', 3: 'width'}
}

torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    dynamic_axes=dynamic_axes,
    opset_version=11
)

然后在TensorRT里，你得用优化配置文件（Optimization Profile）来指定输入的范围：

// 创建优化配置文件
IOptimizationProfile* profile = builder->createOptimizationProfile();

// 设置输入的最小、最优、最大尺寸
profile->setDimensions("input", OptProfileSelector::kMIN, Dims4(1, 3, 192, 192));
profile->setDimensions("input", OptProfileSelector::kOPT, Dims4(1, 3, 224, 224));
profile->setDimensions("input", OptProfileSelector::kMAX, Dims4(1, 3, 256, 256));

// 把配置文件加到config里
config->addOptimizationProfile(profile);

这里有个关键点：最优尺寸（OPT）决定了TensorRT做算子融合和内存优化的基准。你想想看，如果你设的最优尺寸是224x224，但实际跑的是192x192，性能可能会打折扣。我建议你把最常用的尺寸设成OPT。

注意：动态形状不是万能的。有些算子（比如全连接层）要求输入尺寸固定。如果你用了动态形状，这些算子可能会报错。我曾经在一个项目里，模型前半部分是卷积（支持动态），后半部分是全连接（不支持动态），结果折腾了两天才找到问题。

4.4 调试ONNX导入失败问题

ONNX导入失败，是每个TensorRT开发者都会遇到的噩梦。我总结了几种常见情况：

错误类型 典型表现 解决方法
算子不支持 日志里出现"Unsupported operator XXX" 写插件，或者用ONNX-Simplifier简化模型
形状不匹配 "Tensor shape mismatch" 检查ONNX导出时的动态轴设置
版本冲突 "ONNX opset version too high" 降低opset版本，一般用11或13比较稳
常量折叠失败 "Constant folding failed" 用onnxruntime跑一遍，看能不能正常推理

调试时，我一般会这么做：

  1. 看日志：把日志级别调到kVERBOSE，看看到底卡在哪个节点上。
  2. 用Netron可视化：把ONNX文件拖进Netron，看看结构对不对。我经常发现导出的ONNX里多了一些莫名其妙的节点。
  3. 分段测试：把模型切成几段，分别导入。这样能快速定位问题出在哪一段。
  4. 用onnxruntime验证：先用onnxruntime跑一遍ONNX，确认模型本身没问题。如果onnxruntime都跑不通，那肯定是ONNX导出有问题。

避坑指南：我曾经遇到过一个特别诡异的问题——ONNX解析成功，但推理结果全是错的。查了两天，最后发现是输入数据的预处理跟训练时不一致。说白了，ONNX只负责网络结构，不负责数据预处理。你得自己保证输入数据跟训练时一模一样。

嗯，这一章的内容差不多就这些。ONNX解析看起来简单，但实际坑不少。我的建议是：先跑通一个简单的模型，比如一个只有卷积和ReLU的网络，然后再逐步增加复杂度。这样出了问题，你至少知道是哪个新加的算子搞的鬼。

下一章咱们聊聊插件开发的核心——如何写一个自定义算子。到时候我会拿一个实际案例来演示，保证你能跟着做出来。

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