Pytorch模型转onnx模型实例

yipeiwu_com6年前Python基础

如下所示:

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

原因是pytorch中的tensor.size(1)方式onnx识别不了,需要修改成常量。

以上这篇Pytorch模型转onnx模型实例就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

python Matplotlib底图中鼠标滑过显示隐藏内容的实例代码

python Matplotlib底图中鼠标滑过显示隐藏内容的实例代码

在使用Matplotlib画图过程中,有些内容必须鼠标点击或者划过才可以显示,这个问题可以依赖于annotate(s='str' ,xy=(x,y) ,xytext=(l1,l2) ,....

简单谈谈python的反射机制

  对编程语言比较熟悉的朋友,应该知道“反射”这个机制。Python作为一门动态语言,当然不会缺少这一重要功能。然而,在网络上却很少见到有详细或者深刻的剖析论文。下面结合一个web路由的...

python定义类self用法实例解析

这篇文章主要介绍了python定义类self用法实例解析,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下 在定义类的过程中,无论是显式的...

Python向日志输出中添加上下文信息

除了传递给日志记录函数的参数(如msg)外,有时候我们还想在日志输出中包含一些额外的上下文信息。比如,在一个网络应用中,可能希望在日志中记录客户端的特定信息,如:远程客户端的IP地址和用...

python魔法方法-属性访问控制详解

python魔法方法-属性访问控制详解

属性访问控制 所谓的属性访问控制就是控制点号访问属性的行为,而且不仅是类的外部,连类的内部也受控制,代码见真章,边看代码边解释: •__getattr__(self, ite...