自定义PyTorch模型转ONNX遇RuntimeError问题及解决方法

在Jetson设备上尝试将自定义PyTorch模型转换为TensorRT并运行时，导出.pt模型为ONNX格式触发如下错误：

RuntimeError: Unsupported: ONNX export of index_put in opset 9. Please try opset version 11.

触发错误的代码如下：

import torch.onnx
# 标准ImageNet输入 - 3通道，224x224
# 输入值不影响网络结构导出，也可以使用真实输入
dummy_input = torch.randn(1, 3, 224, 224)
# 执行导出
# torch.onnx.export(model, dummy_input, "best.onnx")

通过在torch.onnx.export()函数中添加opset_version=11参数即可解决该问题，修改后的代码如下：

import torch.onnx
# 标准ImageNet输入 - 3通道，224x224
# 输入值不影响网络结构导出，也可以使用真实输入
dummy_input = torch.randn(1, 3, 224, 224)
# 执行导出
torch.onnx.export(model, dummy_input, "best.onnx", opset_version=11)

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