pytorch转onnx报错： Failed to export an ONNX attribute ‘onnx::Gather’, since it’s not constant

pytorch转onnx报错， Failed to export an ONNX attribute ‘onnx::Gather’, since it’s not constant, please try to make things (e.g., kernel size) static if possible

报错信息

RuntimeError: Failed to export an ONNX attribute ‘onnx::Gather’, since it’s not constant, please try to make things (e.g., kernel size) static if possible

直接从报错信息上无法定位到具体的算子

定位方法

通过查看详细的错误信息

  File "/home/ubuntu/anaconda3/envs/profiler/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py", line 167, in wrapper
    args = [_parse_arg(arg, arg_desc, arg_name, fn_name)  # type: ignore[assignment]
  File "/home/ubuntu/anaconda3/envs/profiler/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py", line 167, in <listcomp>
    args = [_parse_arg(arg, arg_desc, arg_name, fn_name)  # type: ignore[assignment]
  File "/home/ubuntu/anaconda3/envs/profiler/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py", line 84, in _parse_arg
    raise RuntimeError("Failed to export an ONNX attribute '" + v.node().kind() +

发现异常在“torch/onnx/symbolic_helper.py", line 84”抛出

在torch源码的该位置增加print信息：

print(v.node())

重新运行转模型代码，发现报错信息多了实际引起该错误信息的代码位置。

%37 : Long(device=cpu) = onnx::Gather[axis=0](%34, %36) # /×××/pytorch_profiler/torch2onnx.py:18:0
        
        shape = x.shape()[1:]

解决方法

分析原因，是tensor的size()或shape方法无法被onnx识别，需要修改为常量。

修改方式有很多，比如使用List结果报错tensor的shape信息：

shape = [int(e) for e in x.shape[1:]]

修改后，再次运行转模型代码，不会出现该错误