ubuntu18 pytorch模型转onnx

转换

torch.onnx.export(model, args, f, export_params=True, verbose=False, input_names=None, 
output_names=None,do_constant_folding=True,dynamic_axes=None,opset_version=9)

常用参数：
    1.model:torch.nn.model
        要导出的模型
    2.args:tuple or tensor
        模型的输入参数。注意tuple的最后参数为dict要小心，详见pytorch文档。
        输入参数只需满足shape正确，为什么要输入参数呢？因为后面torch.jit.trace要用到，先按下不表。
    3.f:file object or string
        转换输出的模型的位置，如'yolov4.onnx'
    4.export_params:bool,default=True
        true表示导出trained model，否则untrained model。默认即可
    5.verbose:bool,default=False
        true表示打印调试信息
    6.input_names:list of string，default=None
        指定输入节点名称
    7.output_names:list of string，default=None
        指定输出节点名称
    8.do_constant_folding:bool,default=True
        是否使用常量折叠，默认即可
    9.dynamic_axes:dict<string, dict<int, string>> or dict<string, list(int)>,default=None
        有时模型的输入输出是可变的，如RNN，或者输入输出图片的batch是可变的，
        这时我们通过dynamic_axes来指定输入tensor的哪些参数可变。
        如input的shape为(b,3,h,w)其中b、h、w可变
          output的shape为(b,c,5)其中b、c可变
        有几种方式实现：
        1)仅list(int)
            dynamic_axes={'input':[0,2,3],'output':[0,1]}
        2)仅dict<int, string>
            dynamic_axes={'input':{0:'batch',2:'height',3:'width'},'output':{0:'batch',1:'c'}}
        3)mixed
            dynamic_axes={'input':{0:'batch',2:'height',3:'width'},'output':[0,1]}
     10.opset_version:int,default=9
        指定onnx的opset版本，版本过低的话，不支持upsample等操作。

例子：

1.单输入输出

import torch
import torch.nn as nn
import onnx
import numpy as np
class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1=nn.Conv2d(3,3, kernel_size=3, stride=2,padding=1)
    def forward(self,x,y):
        result1=self.conv1(x)
        result2=self.conv1(y)
        return result1,result2

model=Model()
model.eval() # 若存在batchnorm、dropout层则一定要eval()!!!!再export

input_names = ["input"]
output_names = ["output"]

input=torch.randn((1,3,12,12))

# 单输入单输出
torch.onnx.export(model,input,'model.onnx',export_params=True, verbose=False, opset_version=10,input_names=input_names,output_names=output_names)
# 如果需要指定input和output的batch可变，可以添加dynamic_axes，如下：
#torch.onnx.export(model,input,'model.onnx',export_params=True, verbose=False, opset_version=10,input_names=input_names,output_names=output_names,dynamic_axes={'input':[0],'output':[0]} )

2.多输入输出

import torch
import torch.nn as nn
import onnx
import numpy as np
class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1=nn.Conv2d(3,3, kernel_size=3, stride=2,padding=1)
    def forward(self,x,y):
        result1=self.conv1(x)
        result2=self.conv1(y)
        return result1,result2

model=Model()
model.eval() # 若存在batchnorm、dropout层则一定要eval()!!!!再export

input_names = ["input_0","input_1","input2"]
output_names = ["output_0","output_1"]

input0=torch.randn((1,3,12,12))
input1=torch.randn((1,3,6,6))
input2=torch.randn(3)

# 多输入使用(input0,input1,input2)方式会报错，改成[input0,input1,input2]
#torch.onnx.export(model,(input0,input1,input2),'model.onnx',export_params=True, verbose=False, opset_version=10,input_names=input_names,output_names=output_names)
torch.onnx.export(model,[input0,input1,input2],'model.onnx',export_params=True, verbose=False, opset_version=10,input_names=input_names,output_names=output_names)
# 如果需要指定input_0和output_0的batch可变，可以添加dynamic_axes，如下：
#torch.onnx.export(model,[input0,input1,input2],'model.onnx',export_params=True, verbose=False, opset_version=10,input_names=input_names,output_names=output_names,dynamic_axes={'input_0':[0],'output_0':[0]} )

参考：

https://zhuanlan.zhihu.com/p/422290231