3、ONNX算子兼容性：常见PyTorch算子到ONNX的映射，遇到不兼容算子怎么办。
好，咱们继续往下走。
上一章我们把模型从 PyTorch 转成了 ONNX，跑通了第一个流程。但说实话，那只是理想情况。实际项目中，你大概率会遇到一个头疼的问题——算子不兼容。
什么叫不兼容？就是 PyTorch 里好好的一个操作，转到 ONNX 时，它不认识，或者映射错了。轻则报错，重则导出成功但推理结果完全不对。
我刚开始做部署时，就被这个坑过。一个简单的 torch.where，导出 ONNX 后，推理结果全是错的。排查了两天，最后发现是 ONNX 的算子集版本太低，不支持这个动态行为。
今天我们就来彻底聊聊这件事。
3.1 常见 PyTorch 算子到 ONNX 的映射关系
首先，你得知道哪些算子是“安全的”。说白了，就是 PyTorch 和 ONNX 之间已经建立好映射的算子。
ONNX 官方维护了一个算子列表，PyTorch 的 torch.onnx 模块会尽量把每个 PyTorch 算子映射到对应的 ONNX 算子。我整理了一份常用映射表，你可以直接参考：
| PyTorch 算子 | ONNX 算子 | 兼容性说明 |
|---|---|---|
torch.nn.Conv2d |
Conv |
完全兼容，注意 padding 和 dilation 参数 |
torch.nn.BatchNorm2d |
BatchNormalization |
完全兼容，训练和推理模式需区分 |
torch.nn.ReLU |
Relu |
完全兼容 |
torch.nn.MaxPool2d |
MaxPool |
完全兼容，注意 ceil_mode |
torch.nn.Linear |
Gemm 或 MatMul |
完全兼容，取决于是否带 bias |
torch.add |
Add |
完全兼容 |
torch.cat |
Concat |
完全兼容，注意 axis 参数 |
torch.reshape |
Reshape |
完全兼容，但动态 shape 时需小心 |
torch.transpose |
Transpose |
完全兼容 |
torch.where |
Where |
部分兼容，ONNX opset < 9 不支持 |
torch.nn.functional.interpolate |
Resize |
部分兼容，不同 opset 行为有差异 |
看到没？大部分基础算子都是没问题的。但像 torch.where、torch.nn.functional.interpolate 这种，就有点“挑版本”了。
3.2 遇到不兼容算子怎么办？
好，重点来了。当你遇到一个 PyTorch 算子，ONNX 就是不认识，或者映射错了，怎么办？
别慌。我有三招，按顺序试，基本能解决 90% 的问题。
3.2.1 第一招：升级 ONNX opset 版本
这是最简单的一招。ONNX 的算子集（opset）是不断更新的。新版本会加入更多算子支持。
比如刚才说的 torch.where，在 opset 9 之前是不支持的。你只需要在导出时指定更高的 opset 版本：
import torch
model = YourModel()
dummy_input = torch.randn(1, 3, 224, 224)
# 指定 opset_version=11 或更高
torch.onnx.export(
model,
dummy_input,
"model.onnx",
opset_version=11, # 试试 11、12、13...
input_names=["input"],
output_names=["output"]
)
我遇到过好几次，升级 opset 版本后，问题就消失了。但要注意，opset 版本越高，对推理引擎的要求也越高。比如 TensorRT 8.x 最高支持 opset 13，你如果设成 15，TensorRT 可能就加载不了。
3.2.2 第二招：用 PyTorch 原生算子替换
如果升级 opset 版本还不行，那就换个思路——用 PyTorch 里已有的、兼容性好的算子，组合实现不兼容算子的功能。
举个例子。假设你的模型里用了一个自定义的 torch.nn.functional.softshrink，这个算子在 ONNX 里没有直接映射。怎么办？
你可以用 torch.where 和 torch.abs 组合实现它：
import torch
import torch.nn as nn
class SoftShrink(nn.Module):
def __init__(self, lambd=0.5):
super().__init__()
self.lambd = lambd
def forward(self, x):
# 用 torch.where 和 torch.abs 模拟 softshrink
return torch.where(
x > self.lambd,
x - self.lambd,
torch.where(
x < -self.lambd,
x + self.lambd,
torch.zeros_like(x)
)
)
# 然后替换模型中的 softshrink
model = YourModel()
model.softshrink = SoftShrink(lambd=0.5)
你看，这样就把一个不兼容的算子，拆解成了两个兼容的算子（where 和 abs）。ONNX 就能识别了。
3.2.3 第三招：自定义 ONNX 算子（终极方案）
如果前两招都不行，那就只能上大招了——自定义 ONNX 算子。
什么意思？就是你自己写一个 ONNX 不认识的算子，然后告诉 ONNX 怎么处理它。
具体做法是：
- 在 PyTorch 中，用
torch.autograd.Function定义一个自定义操作。 - 在导出时，用
torch.onnx.register_custom_op_symbolic注册这个操作到 ONNX 的映射。 - 在推理后端（比如 TensorRT），用插件（Plugin）实现这个算子的推理逻辑。
举个例子。假设你有一个自定义的 MyCustomOp，它做的是 y = x * 2 + 1：
import torch
class MyCustomOp(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
return x * 2 + 1
@staticmethod
def symbolic(g, x):
# 告诉 ONNX：这个操作等价于 "Add(Mul(x, 2), 1)"
return g.op("Add", g.op("Mul", x, g.op("Constant", value_t=torch.tensor(2))),
g.op("Constant", value_t=torch.tensor(1)))
# 注册自定义算子
torch.onnx.register_custom_op_symbolic("my_custom_op", MyCustomOp.symbolic, 1)
# 在模型中使用
class MyModel(torch.nn.Module):
def forward(self, x):
return MyCustomOp.apply(x)
这样，ONNX 就能理解你的自定义操作了。但注意，这只是在 ONNX 层面“骗”过去了。真正推理时，你的推理后端（比如 TensorRT）还需要一个对应的插件来实现这个算子的计算逻辑。
3.3 避坑指南：我踩过的几个坑
最后，分享几个我实际项目中踩过的坑，希望能帮你少走弯路。
- 动态 shape 问题： 有些算子（比如
Reshape）在动态 shape 下会导出失败。我的建议是，导出时用固定 shape，推理时再处理动态输入。 - 控制流问题：
if、for循环在 ONNX 里很难处理。尽量用torch.where替代条件判断。 - 算子版本不一致： 同一个算子，在不同 opset 版本下行为可能不同。比如
Resize在 opset 10 和 opset 11 中的坐标模式就不一样。导出前一定要确认。 - 训练和推理模式：
BatchNorm、Dropout在训练和推理时行为不同。导出前记得调用model.eval()。
嗯，今天就先聊这么多。算子兼容性这块，说白了就是“知己知彼”。你了解 PyTorch 和 ONNX 的映射关系，知道怎么排查和解决，就能少踩很多坑。
下一章，我们会进入 TensorRT 的世界。到时候你会发现，ONNX 只是中间人，真正的性能优化才刚刚开始。