PyTorch 2.0.1导出SPIGA模型到ONNX遇算子不支持问题求助

问题概况 #

• 开发环境：PyCharm，PyTorch 2.0.1+cu117，torchvision 0.15.2+cu117
• 操作目标：将SPIGA模型导出为ONNX格式
• 尝试情况：测试了opset版本7到20，均导出失败
• 核心报错：

Exporting the operator 'aten::affine_grid_generator' to ONNX opset version X is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues

导出使用的代码：

# Create three dummy tensors with the specified sizes
input_image = torch.randn(1, 3, 256, 256).cuda()  # Size: (batch_size, channels, height, width)
landmarks = torch.randn(1, 98, 3).cuda()  # Size: (batch_size, num_landmarks, 3)
cam_matrix = torch.randn(1, 3, 3).cuda()  # Size: (batch_size, 3, 3)

# Create a list containing these tensors
dummy_input = [input_image, landmarks, cam_matrix]
#dummy_input = ([torch.randn(1, 3, 256, 256).cuda(), self.model3d, self.cam_matrix])
onnx_model_path = "spiga_model.onnx"  # Output ONNX file path
torch.onnx.export(
    self.model,  # Your SPIGA model instance
    dummy_input,  # Example input data
    onnx_model_path,  # Output ONNX file path
    verbose=True,  # Enable verbose mode for debugging (optional)
    input_names=self.model_inputs,  # List of input names (adjust as needed)
    output_names=["features"],  # List of output names (adjust as needed)
    opset_version=14  # ONNX opset version (adjust as needed)
)

print(f"SPIGA model exported to {onnx_model_path}")

可行解决方案 #

1. 替换模型中的affine_grid调用 #

aten::affine_grid_generator是torch.nn.functional.affine_grid的底层依赖算子，PyTorch 2.0.x版本对其ONNX支持存在缺陷。在SPIGA模型代码中找到调用affine_grid的位置，替换为以下兼容实现：

import torch.nn.functional as F

def onnx_compatible_affine_grid(theta, size):
    # 保持与原算子输出一致，避免触发不支持的底层调用
    grid = F.affine_grid(theta, size, align_corners=False)
    return grid

注意：若原代码使用align_corners=True，需同步修改参数，防止输出结果偏差。

2. 升级PyTorch版本 #

PyTorch 2.1及以上版本修复了affine_grid的ONNX导出支持问题，执行以下命令升级到对应CUDA版本的稳定版：

pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

升级后推荐使用opset_version=17或18重新尝试导出。

3. 注册自定义ONNX算子映射 #

在导出代码前添加自定义算子映射，将aten::affine_grid_generator转换为ONNX支持的逻辑：

from torch.onnx import register_custom_op_symbolic

def affine_grid_generator_symbolic(g, theta, size, align_corners=False):
    # 构造ONNX兼容的网格生成逻辑
    grid = g.op("AffineGridGenerator", theta, size, align_corners_i=int(align_corners))
    return grid

register_custom_op_symbolic('aten::affine_grid_generator', affine_grid_generator_symbolic, opset_version=14)

不同opset版本可能需要调整映射逻辑，需根据实际情况修改。

4. 临时跳过算子检查（仅作测试用） #

在torch.onnx.export中添加参数跳过部分ONNX检查，可能绕过报错，但存在模型不可用风险，仅用于临时验证：

torch.onnx.export(
    self.model,
    dummy_input,
    onnx_model_path,
    verbose=True,
    input_names=self.model_inputs,
    output_names=["features"],
    opset_version=14,
    do_constant_folding=True,
    enable_onnx_checker=False
)

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