时间:2021-05-22
如下所示:
import ioimport torchimport torch.onnxfrom models.C3AEModel import PlainC3AENetCBAMdevice = 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模型实例就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持。
声明:本页内容来源网络,仅供用户参考;我单位不保证亦不表示资料全面及准确无误,也不保证亦不表示这些资料为最新信息,如因任何原因,本网内容或者用户因倚赖本网内容造成任何损失或损害,我单位将不会负任何法律责任。如涉及版权问题,请提交至online#300.cn邮箱联系删除。
1Pytorch以ONNX方式保存模型defsaveONNX(model,filepath):'''保存ONNX模型:parammodel:神经网络模型:par
目标是想把在服务器上用pytorch训练好的模型转换为可以在移动端运行的tflite模型。最直接的思路是想把pytorch模型转换为tensorflow的模型,
有一些非常流行的网络如resnet、squeezenet、densenet等在pytorch里面都有,包括网络结构和训练好的模型。pytorch自带模型网址:h
1模型定义 和TF很像,Pytorch也通过继承父类来搭建模型,同样也是实现两个方法。在TF中是__init__()和call(),在Pytorch中则是__
今天用pytorch保存模型时遇到bugCan'tpickle在google上查找原因,发现是保存时保存了整个模型的原因,而模型中有一些自定义的参数将torch