Pytorch——初探onnx（1）解決upsample_bilinear2d轉換問題

問題

將pytorch的pth文件讀取后轉換為ONNX格式，嘗試使用下述代碼將model導出為onnx格式

onnx.export(model, input_image, 'model.onnx')

出現了如下錯誤
RuntimeError: ONNX export failed: Couldn‘t export operator aten::upsample_bilinear2d

環境

目前測試發現pytorch1.6及以下均存在此問題

復現

下述代碼中使用的mode為bilinear，align_corners設置為True

nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

解析

在pytorch轉onnx的上采樣操作中的align_corners設置為True之后，就為upsample_bilinear2d操作了。將align_corners=False對應到onnx::Upsample

解決方案一：

替換為onnx支持的操作，調用 torch.nn.functional.interpolate 時（即在模型定義處修改），將 mode 從 bilinear 改成默認的 nearest，注意這個模式下不能設align_corners。改過后可以轉onnx，但會丟精度

解決方案二：

# 報錯代碼
torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names, output_names=output_names)
# 修改增加參數：opset_version=11
torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names, output_names=output_names, opset_version=11)

參考

https://blog.csdn.net/u011622208/article/details/100779248
https://blog.csdn.net/weixin_40437821/article/details/113633166