pytorch导出onnx简单易懂（固定输入尺寸以及动态输入尺寸）

先简单定义一个网络结构

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.ClassifierLayer = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=4, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=4, out_channels=4, kernel_size=3, stride=1),
            nn.MaxPool2d(kernel_size=2),
        )


    def forward(self, x):
        y = self.ClassifierLayer(x)
        y = F.interpolate(y,size=x.shape[2:])
        return y

转换代码（固定输入尺寸）

#实例化网络结构
model = Net()
#定义一个模拟的输入
input_data = torch.randn([1,3,320,320])
#参数1：模型  参数2：模拟的输入数据  参数3：导出的onnx绝对路径  #参数4：是否打印导出过程中的详细信息
torch.onnx.export(model, input_data, "Net.onnx", verbose=True)

接下来可视化一下导出后的onnx网络结构

注意看右边的输入部分[1,3,320,320]

动态输入尺寸转换代码

    input_name = 'input'
    output_name = 'output'
    torch.onnx.export(model, 
                      input_data, 
                      "simplenet.onnx",
                      opset_version=11,
                      input_names=[input_name],
                      output_names=[output_name],
                      dynamic_axes={
                          input_name: {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
                          output_name: {0: 'batch_size', 2: 'out_width', 3: 'out_height'}}
                      )

同样注意右边的输入部分，现在变成了[batch_size,3,in_width,int_height]

至此最基础的pytorch转onnx就全部完成了，所有的网络结构都可套用这两段代码，不过如果遇到了转换失败算子不支持的问题，可以尝试修改opset_version参数，或者使用其他结构去代替不支持的操作