ONNX模型转换过程中算子兼容性报错与动态形状支持问题的从PyTorch到TensorRT的完整坑点排查与修复指南

一、踩坑前先搞懂：为啥要转ONNX再转TensorRT？

很多做模型部署的人都绕不开这两步，先得说清楚为啥要这么干，不然踩坑都不知道为啥踩。先讲应用场景：比如你用PyTorch训了个分类模型，想放到手机、边缘设备或者服务器上跑，要求速度快、占内存小，这时候就需要转成TensorRT的格式——因为TensorRT是英伟达专门搞推理加速的工具，能把模型剪枝、量化，还能优化算子组合，速度能快好几倍甚至十几倍。但TensorRT不能直接读PyTorch的模型，得先转成中间格式ONNX，相当于“翻译”的中间站，ONNX的定位是跨框架的统一模型格式，所有主流框架（PyTorch、TensorFlow、MXNet）都能转成它，然后再转成TensorRT能用的格式。

再讲这个流程的优缺点：优点是能打通不同框架的模型，拿到TensorRT的加速效果；缺点就是中间多了一步，容易出问题，最常见的就是算子不兼容、动态形状支持差，这也是咱们要讲的核心坑。注意事项：转之前一定要把PyTorch模型的推理逻辑固定好，比如有没有用动态输入、有没有自定义算子，这些都是后面出问题的源头。

二、第一大坑：算子兼容性报错的排查与修复

这是最常见的问题，说白了就是PyTorch里的某个算子，ONNX不认识，或者ONNX认识但TensorRT不支持，转的时候就会报错。

2.1 先搞懂报错的本质

举个例子，你转的时候可能会看到这样的错误：“Could not export ONNX model: Unsupported operator: Torch::MyCustomOp”，或者转完ONNX用TensorRT加载的时候报错“Unsupported ONNX operator: 'MyCustomOp'”。本质就是“翻译”环节丢了东西，或者中间站没这个词，或者最终站不认这个词。

2.2 排查步骤：从PyTorch到ONNX再到TensorRT层层查

2.2.1 第一步：确认PyTorch转ONNX时的算子报错

先拿一个完整的例子来，技术栈统一用PyTorch 2.0.0、ONNX 1.14.0、TensorRT 8.6.1，先写一个有问题的PyTorch模型，故意用一个ONNX不支持的算子：

# 技术栈：PyTorch 2.0.0, ONNX 1.14.0
import torch
import torch.nn as nn
import torch.onnx

# 自定义一个PyTorch算子，比如自己写的图像裁剪函数（只是例子，实际可能是更复杂的自定义逻辑）
class MyCropOp(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x, crop_size):
        # 只保留输入张量的前crop_size个通道
        return x[:, :crop_size, :, :]

# 完整的分类模型
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.crop = MyCropOp()  # 引入自定义算子
        self.fc = nn.Linear(64*224*224, 10)  # 假设输入是224x224的图

    def forward(self, x, crop_size):
        x = self.conv1(x)
        x = self.crop(x, crop_size)
        x = x.view(x.size(0), -1)
        return self.fc(x)

# 初始化模型和输入
model = MyModel()
model.eval()  # 推理模式必须开
# 输入是动态的，先给一个样例输入（ONNX转的时候需要样例输入来推导形状）
dummy_input = torch.randn(1, 3, 224, 224)
dummy_crop_size = torch.tensor(32)  # 样例裁剪大小

# 尝试转ONNX
try:
    torch.onnx.export(
        model,
        (dummy_input, dummy_crop_size),  # 样例输入，多输入要放元组
        "my_model.onnx",
        opset_version=12,  # ONNX的算子版本，越高支持的算子越多
        input_names=["input", "crop_size"],
        output_names=["output"],
        dynamic_axes={  # 先定义动态轴，后面讲
            "input": {0: "batch_size"},
            "crop_size": {0: "batch_size"}
        }
    )
except Exception as e:
    print("转ONNX报错：", e)

运行这段代码，肯定会报错，因为MyCropOp是PyTorch的自定义算子，ONNX根本不认识。

2.2.2 第二步：修复算子的方法

有几种修复方法，按优先级来：

1. 替换成ONNX支持的原生算子：这是最好的方法，因为原生算子兼容性最好。比如上面的MyCropOp，其实可以用PyTorch的切片操作代替，把MyCropOp改成：

# 技术栈：PyTorch 2.0.0, ONNX 1.14.0
class MyCropOp(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x, crop_size):
        # 用PyTorch的切片，这个操作ONNX是支持的
        return x[:, :crop_size, :, :]

再转ONNX就不会报错了，因为切片是PyTorch的原生操作，ONNX有对应的Slice算子。

2. 用ONNX支持的自定义算子（ONNX Custom Op）：如果实在不能替换，比如你的自定义算子有特殊逻辑，那可以用ONNX的自定义算子。步骤是：先在PyTorch里给自定义算子加装饰器，告诉ONNX这个算子的定义，然后转ONNX，最后转TensorRT的时候再实现这个算子。比如给MyCropOp加装饰器：

# 技术栈：PyTorch 2.0.0, ONNX 1.14.0
from torch.onnx import register_custom_op_symbolic

# 定义ONNX端的算子符号，告诉ONNX这个算子叫什么，输入输出是什么
def my_crop_symbolic(g, x, crop_size):
    # g是ONNX的计算图对象，返回的是ONNX的算子节点
    return g.op("my_custom::MyCrop", x, crop_size)

# 把PyTorch的算子和ONNX的算子符号绑定
register_custom_op_symbolic("::MyCrop", my_crop_symbolic, opset_version=12)

# 然后再定义MyCropOp的时候，要给PyTorch的算子注册
class MyCropOp(nn.Module):
    def __init__(self):
        super().__init__()
        # 注册PyTorch的自定义算子，名字要和上面的一致
        self.crop = torch.nn.functional.relu  # 这里只是占位，实际是自定义逻辑
        # 正确的注册方式（PyTorch 2.0的例子）
        from torch._ops import OpOverload
        class MyCropOp(OpOverload):
            @staticmethod
            def forward(ctx, x, crop_size):
                return x[:, :crop_size, :, :]
            @staticmethod
            def symbolic(g, x, crop_size):
                return g.op("my_custom::MyCrop", x, crop_size)
        self.crop = MyCropOp

    def forward(self, x, crop_size):
        return self.crop(x, crop_size)

转完ONNX后，这个算子就会变成my_custom::MyCrop，然后转TensorRT的时候，需要自己写这个算子的CUDA实现，再注册到TensorRT里，这个比较复杂，适合实在没办法的情况。

3. 把算子放到推理后处理：如果算子的逻辑很简单，比如只是一个简单的计算，那可以把它从模型里拿出来，放到推理的后处理步骤，比如用Python或者C++实现，这样模型里就没有这个算子了，转的时候就不会报错。比如上面的裁剪操作，完全可以在推理完模型后，用Python把输出的张量裁剪一下，这样模型里就不需要MyCropOp了。

2.2.3 第三步：确认ONNX转TensorRT时的算子报错

有时候转ONNX没问题，但转TensorRT的时候会报错，因为TensorRT支持的ONNX算子比ONNX本身少，比如ONNX有某个算子，但TensorRT不支持。这时候怎么排查？可以用ONNX的工具检查算子：

# 技术栈：ONNX 1.14.0
# 安装onnxruntime工具
pip install onnxruntime
# 检查ONNX模型的算子，看有没有不支持的
python -m onnxruntime.tools.convert_onnx_models_to_ort my_model.onnx --optimization_level all

如果有不支持的算子，会输出“Unsupported operator: X”，然后再按上面的方法替换。

三、第二大坑：动态形状支持的排查与修复

动态形状是什么？比如你训练的时候输入是224x224的图，但推理的时候可能输入是256x256，或者批次大小是1、2、3都有可能，这就是动态形状。很多人转的时候会遇到“输入形状不匹配”的错误，本质就是转的时候没处理好动态形状。

3.1 动态形状的报错本质

举个例子，你转完ONNX，用TensorRT加载的时候报错“Input shape mismatch: expected [1,3,224,224], got [1,3,256,256]”，或者转ONNX的时候没定义动态轴，导致转出来的模型输入形状是固定的。

3.2 排查步骤：从PyTorch转ONNX时的动态轴定义开始

3.2.1 第一步：正确定义ONNX的动态轴

转ONNX的时候有个参数叫dynamic_axes，就是用来定义哪些轴是动态的。很多人写错这个参数，比如定义的轴不对，或者名字不对。比如之前的例子，我们要让输入的批次大小（第0轴）、输入的高度（第2轴）、宽度（第3轴）都是动态的，那dynamic_axes应该这么写：

# 技术栈：PyTorch 2.0.0, ONNX 1.14.0
# 转ONNX时的dynamic_axes参数
dynamic_axes={
    # 输入的第0轴（批次大小）是动态的，名字叫batch_size
    "input": {0: "batch_size", 2: "height", 3: "width"},
    # 裁剪大小的第0轴（和批次大小对应）是动态的
    "crop_size": {0: "batch_size"},
    # 输出的第0轴（批次大小）是动态的
    "output": {0: "batch_size"}
}

这里要注意：每个动态轴的名字要对应，比如输入的第0轴和裁剪大小的第0轴都叫batch_size，意思是这两个轴的大小是一样的，比如批次大小是2，那输入的第0轴是2，裁剪大小的第0轴也是2，不能一个叫batch1一个叫batch2，不然转的时候会出错。

3.2.2 第二步：用样例输入验证动态形状

转ONNX的时候，样例输入的形状可以是任意的，只要是合法的就行，比如你要支持动态的高度和宽度，样例输入可以是(1,3,224,224)，转完之后可以用ONNX Runtime测试一下动态形状：

# 技术栈：ONNX Runtime 1.14.0
import onnxruntime as ort
import numpy as np

# 加载ONNX模型
session = ort.InferenceSession("my_model.onnx")
# 定义动态输入，比如批次大小是2，高度是256，宽度是256
input_data = np.random.randn(2, 3, 256, 256).astype(np.float32)
crop_size_data = np.array([32, 32]).astype(np.int64)  # 批次大小是2，所以裁剪大小是2个
# 推理
output = session.run(["output"], {"input": input_data, "crop_size": crop_size_data})
print("动态形状推理成功，输出形状：", output[0].shape)

如果运行这段代码没有报错，说明ONNX的动态形状定义是对的。

3.2.3 第三步：转TensorRT时的动态形状配置

转TensorRT的时候，不能直接用转ONNX的动态轴，还要配置TensorRT的动态形状范围，比如你要支持的批次大小是1到16，高度是224到512，宽度是224到512，那转的时候要指定这个范围。比如用trtexec工具转：

# 技术栈：TensorRT 8.6.1
# trtexec是TensorRT自带的转模型工具
trtexec --onnx=my_model.onnx --saveEngine=my_model.trt \
  # 定义输入的最小形状、最优形状、最大形状
  --minShapes=input:1x3x224x224,crop_size:1x1 \
  --optShapes=input:8x3x256x256,crop_size:8x1 \
  --maxShapes=input:16x3x512x512,crop_size:16x1

这里要注意：minShapes、optShapes、maxShapes的形状必须符合你在ONNX里定义的动态轴，比如你在ONNX里定义了input的第0轴是动态的，那这三个形状的第0轴必须分别是1、8、16，不能是固定的1。如果转的时候报错“Dynamic shape range is invalid”，说明你配置的范围不对，比如maxShapes的形状超过了模型支持的范围，或者minShapes小于模型支持的最小形状。

3.2.4 第四步：TensorRT推理时的动态形状匹配

转完TensorRT模型后，推理的时候要注意，输入的形状必须在你配置的minShapes和maxShapes之间，比如你配置的maxShapes是16x3x512x512，那你不能输入17x3x512x512的形状，不然会报错。推理的时候还要指定形状，比如用Python加载TensorRT模型：

# 技术栈：TensorRT 8.6.1, PyCUDA 2023.1
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

# 加载TensorRT引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with open("my_model.trt", "rb") as f:
    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()

# 定义输入形状，比如批次大小是2，高度是256，宽度是256
input_shape = (2, 3, 256, 256)
crop_size_shape = (2, 1)
# 设置输入的形状，必须在min和max之间
context.set_input_shape("input", input_shape)
context.set_input_shape("crop_size", crop_size_shape)

# 准备输入数据
input_data = np.random.randn(*input_shape).astype(np.float32)
crop_size_data = np.array([32, 32]).astype(np.int64).reshape(crop_size_shape)
# 分配CUDA内存
d_input = cuda.mem_alloc(input_data.nbytes)
d_crop = cuda.mem_alloc(crop_size_data.nbytes)
# 把数据拷贝到CUDA内存
cuda.memcpy_htod(d_input, input_data)
cuda.memcpy_htod(d_crop, crop_size_data)

# 准备输出
output_shape = (2, 10)
d_output = cuda.mem_alloc(np.prod(output_shape) * np.float32().nbytes)
# 绑定输入输出
context.set_tensor_address("input", int(d_input))
context.set_tensor_address("crop_size", int(d_crop))
context.set_tensor_address("output", int(d_output))

# 推理
context.execute_v2([int(d_input), int(d_crop), int(d_output)])

# 把输出拷贝到CPU
output = np.empty(output_shape, dtype=np.float32)
cuda.memcpy_dtoh(output, d_output)
print("TensorRT动态形状推理成功，输出：", output)

四、常见坑的汇总与避坑指南

除了上面的两个大坑，还有一些小坑，比如：

1. PyTorch模型的推理模式没开：转ONNX的时候必须开model.eval()，不然模型里的BN层、Dropout层会用训练模式的逻辑，转出来的模型推理结果不对。
2. ONNX的opset版本太低：opset版本越高，支持的算子越多，建议用opset_version=12以上，现在最新的是17，只要你的TensorRT支持就行。
3. TensorRT的版本不匹配：转ONNX用的opset版本要和TensorRT支持的版本匹配，比如TensorRT 8.6支持opset_version=17，TensorRT 8.2支持opset_version=14，版本不匹配会导致算子不支持。
4. 动态形状的轴名字不对应：转ONNX的时候dynamic_axes里的轴名字必须对应，比如输入的第0轴和输出的第0轴都叫batch_size，不能随便取名字。

五、应用场景总结

这个流程的应用场景主要是：

1. 边缘设备部署：比如手机、智能摄像头、自动驾驶的边缘设备，需要模型速度快、占内存小，用TensorRT加速。
2. 服务器推理服务：比如线上的图像分类、目标检测服务，需要高吞吐量，用TensorRT加速能提升服务的并发量。
3. 跨框架部署：比如用PyTorch训的模型，要部署到支持TensorRT的平台，必须转ONNX再转TensorRT。

六、优缺点与注意事项总结

6.1 优缺点

优点：

1. 能拿到TensorRT的加速效果，速度提升明显。
2. 能打通不同框架的模型，实现跨框架部署。
3. 支持动态形状，能适应不同的输入需求。

缺点：

1. 中间多了一步ONNX转译，容易出算子不兼容的问题。
2. 动态形状的配置比较复杂，容易出错。
3. 自定义算子的修复比较麻烦，需要写CUDA代码。

6.2 注意事项

1. 转之前一定要把PyTorch模型的推理逻辑固定好，避免用训练模式的算子。
2. 转ONNX的时候一定要正确定义dynamic_axes，轴名字要对应。
3. 转TensorRT的时候一定要正确配置minShapes、optShapes、maxShapes，范围要符合需求。
4. 转完之后一定要测试，比如用ONNX Runtime测试ONNX模型的正确性，用TensorRT测试动态形状的正确性。

七、文章总结

从PyTorch转ONNX再转TensorRT的流程，核心坑就是算子兼容性和动态形状支持，排查的时候要从PyTorch转ONNX开始，一步步来，先解决转ONNX的问题，再解决转TensorRT的问题。修复算子的优先级是替换成原生算子>放到后处理>自定义算子，动态形状的核心是正确定义轴和配置范围。只要掌握了这些方法，就能顺利完成模型的转换和部署。