4、ONNX导出常见错误：算子不支持、动态轴问题、输入输出命名不一致

说实话，ONNX导出这块，我踩过的坑比吃过的盐还多。很多同学模型训练得好好的，一到导出就翻车。今天咱们就把这几个最要命的错误掰开揉碎了讲清楚。

4.1 算子不支持：最让人头疼的兼容性问题

什么叫算子不支持？说白了就是PyTorch里有的操作，ONNX不认识。我刚开始做部署时，经常遇到导出报错，一看日志——"Unsupported operator"。嗯，这时候别慌，有套路。

常见不支持算子清单：
  • torch.einsum — 这个基本没戏，得手动拆解
  • F.upsample_bilinear2d — 老版本ONNX不支持，建议用torch.nn.functional.interpolate
  • torch.where 的条件分支 — 动态图导出时容易翻车
  • 自定义的 torch.autograd.Function — 这个得自己写符号函数

我在项目中遇到过最离谱的一次，是用了torch.topk，结果ONNX导出时直接报错。后来发现是ONNX opset版本太低，升级到13就解决了。你想想看，有时候问题就这么简单。

4.1.1 解决方案：符号函数注册

对于自定义算子，我们需要手动注册符号函数。我个人习惯的做法是这样的：

import torch.onnx.symbolic_helper as sym_help
from torch.onnx.symbolic_registry import register_op

def my_custom_op_symbolic(g, input, param1, param2):
    """自定义算子的符号函数"""
    # 这里实现ONNX中对应的操作
    return g.op("MyCustomOp", input, 
                param1_i=param1,
                param2_f=param2)

# 注册算子
register_op('my_custom_op', my_custom_op_symbolic, '', 13)
我的经验：注册符号函数时，opset版本一定要和导出时保持一致。我曾经因为版本号写错，折腾了一整天才发现。

4.2 动态轴问题：静态图里的"活"变量

ONNX本质上是静态图，每个张量的形状在导出时就固定了。但实际推理时，batch size可能是变化的，输入图片尺寸也可能不同。这就引出了动态轴的问题。

为什么会这样？因为PyTorch是动态图，而ONNX是静态图。说白了，一个灵活，一个死板。我们需要在两者之间找到平衡。

4.2.1 动态轴的配置方法

导出时，通过dynamic_axes参数来指定哪些维度是可变的：

import torch

# 假设模型输入形状为 [batch, 3, 224, 224]
dummy_input = torch.randn(1, 3, 224, 224)

# 定义动态轴
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes,
    opset_version=13
)
注意：动态轴不是万能的。有些算子对动态形状支持不好，比如torch.reshape如果依赖动态形状，导出后可能出问题。我建议能固定就固定，实在不行再用动态轴。

4.2.2 动态轴的实际坑点

我记得有一次，模型导出时设置了动态batch，结果TensorRT推理时速度反而变慢了。后来发现是动态轴导致TensorRT做了很多优化预留，反而影响了性能。嗯，这里要注意：动态轴虽然灵活，但会牺牲一些优化空间。

场景 建议 原因
固定batch size 不用动态轴 性能最优，TensorRT可以做极致优化
batch size变化范围小 使用动态轴，但限制范围 比如只允许1,2,4,8这几个值
输入尺寸变化大 使用动态轴 但要做好性能测试

4.3 输入输出命名不一致：最隐蔽的错误

这个错误特别坑。模型导出时一切正常，但到了TensorRT推理阶段，输入输出名字对不上，直接报错。我曾经因为这个bug排查了整整一个下午。

说白了，就是导出时指定的名字和推理时用的名字不一致。比如导出时叫input，推理时却用了data。

4.3.1 命名规范建议

我个人习惯的做法是：

  • 输入命名：统一用input、input1、input2这种简单名字
  • 输出命名：统一用output、output1、output2
  • 避免特殊字符：不要用空格、中文、特殊符号
  • 大小写敏感：ONNX和TensorRT都是大小写敏感的，Input和input不一样
避坑指南：我曾经在导出时用了input.1这样的名字，结果TensorRT解析时把点号当成了层级分隔符，直接报错。后来统一改成input_1就解决了。

4.3.2 验证命名一致性的方法

导出后，建议立即用以下代码验证：

import onnx

model = onnx.load("model.onnx")

# 检查输入
for input in model.graph.input:
    print(f"Input name: {input.name}")

# 检查输出
for output in model.graph.output:
    print(f"Output name: {output.name}")

# 检查所有节点
for node in model.graph.node:
    print(f"Node: {node.name}, Inputs: {node.input}, Outputs: {node.output}")

你想想看，如果这里名字就对不上，后面TensorRT推理肯定出问题。所以导出后第一件事，就是检查命名。

4.4 综合实战：一个完整的导出示例

最后，我给大家一个完整的导出模板，把上面三个问题都考虑进去：

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)

# 定义动态轴（如果需要）
dynamic_axes = {
    'input': {0: 'batch_size'},
    'output': {0: 'batch_size'}
}

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

print("导出成功！")
最后提醒：导出成功后，建议用onnxruntime先跑一遍推理，确认没问题再交给TensorRT。这一步能帮你省下大量调试时间。

好了，关于ONNX导出的三个常见错误，今天就聊到这里。记住：算子不支持就注册符号函数，动态轴问题要权衡灵活性和性能，命名不一致就统一规范。把这些搞定了，导出这块基本就不会翻车了。

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