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
原因是pytorch中的tensor.size(1)方式onnx识别不了,需要修改成常量。
以上这篇Pytorch模型转onnx模型实例就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持脚本之家。
相关文章
Python图像滤波处理操作示例【基于ImageFilter类】
这篇文章主要介绍了Python图像滤波处理操作,结合实例形式分析了Python基于ImageFilter类实现的滤波处理相关操作技巧,需要的朋友可以参考下2019-01-01Python3 循环语句(for、while、break、range等)
这篇文章主要介绍了Python3 循环语句(for、while、break、range等),大家把下面的文章看完就基本上就可以了解了python的循环实现方式了2017-11-11安装出现:Requirement already satisfied解决办法
最近pip install的时候报错,一大串Requirement already satisfied,所以下面这篇文章主要给大家介绍了关于安装出现:Requirement already satisfied的解决办法,需要的朋友可以参考下2022-08-08python探索之BaseHTTPServer-实现Web服务器介绍
这篇文章主要介绍了python探索之BaseHTTPServer-实现Web服务器介绍,小编觉得还是挺不错的,这里分享给大家,供需要的朋友参考。2017-10-10
最新评论