Pytorch模型转onnx模型实例

所属分类: 脚本专栏 / python 阅读数: 962
收藏 0 赞 0 分享

如下所示:

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实现图像几何变换的方法,实例分析了Python基于Image模块实现图像翻转、旋转、改变大小等操作的相关技巧,非常简单实用,需要的朋友可以参考下
收藏 0 赞 0 分享

Python中的urllib模块使用详解

这篇文章主要介绍了Python中的urllib模块使用详解,是Python入门学习中的基础知识,需要的朋友可以参考下
收藏 0 赞 0 分享

Python的多态性实例分析

这篇文章主要介绍了Python的多态性,以实例形式深入浅出的分析了Python在面向对象编程中多态性的原理与实现方法,需要的朋友可以参考下
收藏 0 赞 0 分享

python生成IP段的方法

这篇文章主要介绍了python生成IP段的方法,涉及Python文件读写及随机数操作的相关技巧,具有一定参考借鉴价值,需要的朋友可以参考下
收藏 0 赞 0 分享

python操作redis的方法

这篇文章主要介绍了python操作redis的方法,包括Python针对redis的连接、设置、获取、删除等常用技巧,具有一定参考借鉴价值,需要的朋友可以参考下
收藏 0 赞 0 分享

python妹子图简单爬虫实例

这篇文章主要介绍了python妹子图简单爬虫,实例分析了Python爬虫程序所涉及的页面源码获取、进度显示、正则匹配等技巧,需要的朋友可以参考下
收藏 0 赞 0 分享

分析用Python脚本关闭文件操作的机制

这篇文章主要介绍了分析用Python脚本关闭文件操作的机制,作者分Python2.x版本和3.x版本两种情况进行了阐述,需要的朋友可以参考下
收藏 0 赞 0 分享

python实现搜索指定目录下文件及文件内搜索指定关键词的方法

这篇文章主要介绍了python实现搜索指定目录下文件及文件内搜索指定关键词的方法,可实现针对文件夹及文件内关键词的搜索功能,需要的朋友可以参考下
收藏 0 赞 0 分享

python中getaddrinfo()基本用法实例分析

这篇文章主要介绍了python中getaddrinfo()基本用法,实例分析了Python中使用getaddrinfo方法进行IP地址解析的基本技巧,需要的朋友可以参考下
收藏 0 赞 0 分享

python查找指定具有相同内容文件的方法

这篇文章主要介绍了python查找指定具有相同内容文件的方法,涉及Python针对文件操作的相关技巧,需要的朋友可以参考下
收藏 0 赞 0 分享
查看更多