PyTorch模型导出ONNX遇torch_scatter::scatter_max不支持问题求解

方案一：注册自定义算子的ONNX符号映射 #

PyTorch允许为第三方算子注册自定义ONNX符号函数，将torch_scatter.scatter_max映射为ONNX支持的算子组合：

import torch
import torch.nn as nn
from torch.onnx import register_custom_op_symbolic
import torch_scatter

# 定义scatter_max的ONNX符号实现
def scatter_max_symbolic(g, input, index, dim=0, out=None):
    # 处理非0维度的转置
    if dim != 0:
        input = g.op("Transpose", input, perm=[dim, 0] + list(range(1, dim)) + list(range(dim+1, input.type().dim())))
        index = g.op("Transpose", index, perm=[dim, 0] + list(range(1, dim)) + list(range(dim+1, index.type().dim())))
    
    # 获取唯一索引及逆映射
    unique_indices, inverse_indices = g.op("Unique", index, sorted_i=1, return_inverse_i=1)
    # 按唯一索引分组取最大值
    grouped_max = g.op("SegmentMax", input, unique_indices)
    # 将结果映射回原索引位置
    result_max = g.op("Gather", grouped_max, inverse_indices, axis=0)
    
    # 转回原维度顺序
    if dim != 0:
        result_max = g.op("Transpose", result_max, perm=[1, 0] + list(range(2, result_max.type().dim())))
    
    # 补充argmax返回值（简化版，若需要精确argmax需补充逻辑）
    result_argmax = g.op("Constant", value_t=torch.zeros_like(index, dtype=torch.long))
    return result_max, result_argmax

# 注册自定义符号函数
register_custom_op_symbolic('torch_scatter::scatter_max', scatter_max_symbolic, opset_version=11)

# 模型定义
class DGCN(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, index):
        return torch_scatter.scatter_max(x, index, dim=0)

# 导出ONNX
model = DGCN()
dummy_x = torch.randn(10, 5)
dummy_index = torch.randint(0, 3, (10,))
torch.onnx.export(model, (dummy_x, dummy_index), "dgcn.onnx", opset_version=11)

方案二：替换为PyTorch原生可导出的等价实现 #

方法1：使用torch.scatter_reduce（PyTorch≥1.12） #

PyTorch 1.12+原生支持scatter_reduce函数，支持'max'操作且可直接导出ONNX：

import torch.nn as nn

class DGCN(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, index):
        # 替换torch_scatter.scatter_max
        max_vals = torch.scatter_reduce(x, dim=0, index=index.unsqueeze(-1), reduce='max', include_self=False)
        # 若需要argmax，需额外计算原索引位置
        argmax_indices = torch.zeros_like(max_vals, dtype=torch.long)
        # 补充argmax逻辑（示例）
        for idx in torch.unique(index):
            mask = index == idx
            _, argmax = x[mask].max(dim=0)
            argmax_indices[idx] = mask.nonzero()[argmax]
        return max_vals, argmax_indices

方法2：原生函数组合实现（兼容低版本PyTorch） #

import torch.nn as nn

def native_scatter_max(x, index, dim=0):
    unique_idx = torch.unique(index)
    max_vals = []
    argmax_indices = []
    for idx in unique_idx:
        mask = index == idx
        slice_x = x[mask] if dim == 0 else x[:, mask]
        max_val, argmax = slice_x.max(dim=dim)
        max_vals.append(max_val)
        # 映射回原张量的索引
        original_argmax = mask.nonzero()[argmax].squeeze()
        argmax_indices.append(original_argmax)
    
    # 整理为与输入匹配的形状
    result_max = torch.zeros_like(x)
    result_argmax = torch.zeros_like(index, dtype=torch.long)
    result_max.index_copy_(dim, unique_idx, torch.stack(max_vals))
    result_argmax.index_copy_(0, unique_idx, torch.tensor(argmax_indices))
    return result_max, result_argmax

class DGCN(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, index):
        return native_scatter_max(x, index, dim=0)

方案三：启用ONNX自定义算子导出 #

若仅在支持自定义算子的环境（如ONNX Runtime）中运行模型，可直接开启自定义算子导出：

torch.onnx.export(
    model,
    (dummy_x, dummy_index),
    "dgcn.onnx",
    opset_version=15,
    enable_custom_ops=True,
    custom_opsets={"torch_scatter": 1}
)

注意：此方式导出的ONNX模型依赖torch_scatter自定义算子，需确保推理环境支持。

内容的提问来源于stack exchange，提问作者Sajjad Aemmi