Pytorch模型转onnx模型实例_Python

Pytorch模型转onnx模型实例

2020-04-21 09:47joey.lei Python

今天小编就为大家分享一篇Pytorch模型转onnx模型实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧

如下所示：

				?

									import io

									import torch

									import torch.onnx

									from models.C3AEModel import PlainC3AENetCBAM

									device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

									def test():

									  model = PlainC3AENetCBAM()

									  pthfile = r'/home/joy/Projects/models/emotion/PlainC3AENet.pth'

									  loaded_model = torch.load(pthfile, map_location='cpu')

									  # try:

									  #   loaded_model.eval()

									  # except AttributeError as error:

									  #   print(error)

									  model.load_state_dict(loaded_model['state_dict'])

									  # model = model.to(device)

									  #data type nchw

									  dummy_input1 = torch.randn(1, 3, 64, 64)

									  # dummy_input2 = torch.randn(1, 3, 64, 64)

									  # dummy_input3 = torch.randn(1, 3, 64, 64)

									  input_names = [ "actual_input_1"]

									  output_names = [ "output1" ]

									  # torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names)

									  torch.onnx.export(model, dummy_input1, "C3AE_emotion.onnx", verbose=True, input_names=input_names, output_names=output_names)

									if __name__ == "__main__":

									 test()

直接将PlainC3AENetCBAM替换成需要转换的模型，然后修改pthfile，输入和onnx模型名字然后执行即可。

注意：上面代码中注释的dummy_input2，dummy_input3，torch.onnx.export对应的是多个输入的例子。

在转换过程中遇到的问题汇总

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

在转换过程中遇到RuntimeError: Failed to export an ONNX attribute, since it's not constant, please try to make things (e.g., kernel size) static if possible的错误。

根据报的错误日志信息打开/home/joy/.tensorflow/venv/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py，在相应位置添加print之后，可以定位到具体哪个op出问题。

例如：

在相应位置添加

				?

									print（v.node()）

输出信息如下：

				?

									%124 : Long() = onnx::Gather[axis=0](%122, %121), scope: PlainC3AENetCBAM/Bottleneck[cbam]/CBAM[cbam]/ChannelGate[ChannelGate] # /home/joy/Projects/models/emotion/WhatsTheemotion/models/cbam.py:46:0