知识复习 - Pytorch导出ONNX模型 动态输入参量

在使用Pytorch导出ONNX模型的时候，经常需要考虑动态的输入参量，比如可变的batch size以及不同尺寸的图片。我们可以通过修改torch.onnx.export函数的dynamic_axes来考虑动态维度的输入参量。这里给一个简单的例子，我们先定义一个卷积神经网络:

import torch
import torch.nn as nn
 
class Model_Net(nn.Module):
    def __init__(self):
        super(Model_Net, self).__init__()
        self.layer1 = nn.Sequential(
 
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
 
            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )
        
    def forward(self, data):
        data = self.layer1(data)
        return data

然后我们用torch.onnx.export()函数动态导出我们定义的卷积神经网络，注意更改dynamic_axes参数，代码如下：

# 创建CNN
model = Model_Net()

# 评估模式
model.eval()

# dummy input, batch_size=1, 3通道，256x256的图像
# 如果导出的ONNX模型没有使用动态输入参量，那么推理时候的输入参量尺寸要和 "dummy input" 一致 
input_data = torch.randn(1, 3, 256, 256)  

torch.onnx.export(model, 
                  input_data, 
                  "Dynamics_InputNet.onnx",
                  opset_version=12,
                  input_names=["input_name"],
                  output_names=["output_name"],
                  dynamic_axes={
                      "input_name": {0: 'batch_size', 2: 'input_height', 3: 'input_width'},
                      "output_name": {0: 'batch_size', 2: 'output_height', 3: 'output_width'}})

dynamic_axes 参数中的 0、2、3 表示相应的维度设置为动态值。最后我们用onnxruntime 来测试导出的卷积神经网络，看看onnx模型是否能适应不同维度的输入，代码如下：

import numpy as np
import onnx
import onnxruntime
 
if __name__ == "__main__":

    # 定义动态的输入
    input_data1 = np.random.rand(8, 3, 512, 512).astype(np.float32)
    input_data2 = np.random.rand(4, 3, 256, 256).astype(np.float32)
    
    # 加载 onnx 模型
    Onnx_file = "./Dynamics_InputNet.onnx"
    Model = onnx.load(Onnx_file)
    onnx.checker.check_model(Model) # 验证Onnx模型是否准确
    
    # 使用 onnxruntime 推理
    model = onnxruntime.InferenceSession(Onnx_file, providers=['CPUExecutionProvider'])
    input_name = model.get_inputs()[0].name
    output_name = model.get_outputs()[0].name
 
    output1 = model.run([output_name], {input_name:input_data1})
    output2 = model.run([output_name], {input_name:input_data2})
 
    print('output1.shape: ', np.squeeze(np.array(output1), 0).shape)
    print('output2.shape: ', np.squeeze(np.array(output2), 0).shape)

代码运行结果为：

我们可以看到ONNX模型成功接收了两个不同维度的输入参量。如果我们事先导出了一个使用静态输入张量的ONNX模型 （i.e., 不设置任何dynamic_axes参数，包括所有维度），推理时就会遇到以下错误：

在之前的blog里，我们介绍了Netron应用。Netron应用不只可以可视化导出的ONNX模型，还可以分析导出的模型是否使用了动态张量做为输入，

注意这里的 input_height, input_width, 和 batch_size。有兴趣的小伙伴可以看一下如果导出一个使用静态张量做为输入的ONNX模型，这里显示的是什么。小伙伴们，学会了吗～

参考文献：

知识复习 - Pytorch导出ONNX模型以及使用Netron进行ONNX模型的可视化-CSDN博客

Pytorch复习笔记--导出Onnx模型为动态输入和静态输入_onnx动态输入-CSDN博客

知识复习 - Pytorch导出ONNX模型 多输入参量-CSDN博客

TensorRT加速推理入门-1：Pytorch转ONNX-CSDN博客