为什么 pytorch 转成 onnx 的时候，需要固定输出形状？

Question

最近在把 pytorch 训练出来的模型，部署到生产环境中去

9.1 使用ONNX进行部署并推理

使用类似下面的 demo 代码，可以把 pytorch 的模型转成 onnx 的模型

import torch.onnx 
# 转换的onnx格式的名称，文件后缀需为.onnx
onnx_file_name = "xxxxxx.onnx"
# 我们需要转换的模型，将torch_model设置为自己的模型
model = torch_model
# 加载权重，将model.pth转换为自己的模型权重
# 如果模型的权重是使用多卡训练出来，我们需要去除权重中多的module. 具体操作可以见5.4节
model = model.load_state_dict(torch.load("model.pth"))
# 导出模型前，必须调用model.eval()或者model.train(False)
model.eval()
# dummy_input就是一个输入的实例，仅提供输入shape、type等信息 
batch_size = 1 # 随机的取值，当设置dynamic_axes后影响不大
dummy_input = torch.randn(batch_size, 1, 224, 224, requires_grad=True) 
# 这组输入对应的模型输出
output = model(dummy_input)
# 导出模型
torch.onnx.export(model,        # 模型的名称
                  dummy_input,   # 一组实例化输入
                  onnx_file_name,   # 文件保存路径/名称
                  export_params=True,        #  如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False.
                  opset_version=10,          # ONNX 算子集的版本，当前已更新到15
                  do_constant_folding=True,  # 是否执行常量折叠优化
                  input_names = ['input'],   # 输入模型的张量的名称
                  output_names = ['output'], # 输出模型的张量的名称
                  # dynamic_axes将batch_size的维度指定为动态，
                  # 后续进行推理的数据可以与导出的dummy_input的batch_size不同
                  dynamic_axes={'input' : {0 : 'batch_size'},    
                                'output' : {0 : 'batch_size'}})

但是我有个疑问，为什么模型转换的时候，需要指定 batch_size 和 input_shape 这些东西呢？这些东西是实际运行的时候，会根据 CPU、内存大小、GPU、GPU 显存大小等等做出调整的，而不是固定死的才对

Answer 1

                  dynamic_axes={'input' : {0 : 'batch_size'},    
                                'output' : {0 : 'batch_size'}})

这个参数叫 dynamic_axes ，dynamic 就是动态的。只有在这里指定过的，才是可变。

如果不指定，那就只能是固定值。比如你的输入是 1x1x224x224 ，那么导出的 onnx 就只能接受 1x1x224x224 输入。2x1x224x224 就不行。

指定过的，比如想例子里的，输入的第 0 维是 batch_size ，那么输入的维度就是 batch_sizex1x224x224 。于是生成的onnx就既可以接受 1x1x224x224 输入，也可以接受 2x1x224x224 的输入了。

就是说，只有被指定过的维度才是后续可以调整的。没有指定过的没有办法调整的。

另外说一句，上面的 "batch_size" 其实可以是任意字符串。叫 "batch_size" 只是为了可以方便理解模型。

Answer 2

在 PyTorch 中，神经网络的输出形状通常是根据输入形状来自动计算的，而在 ONNX 中，输出形状需要在转换时进行显式指定，这是由于 ONNX 的静态图执行模型与 PyTorch 的动态图执行模型不同所致。

当你将 PyTorch 模型转换为 ONNX 模型时，你需要为 ONNX 模型中的每个输出定义固定的形状，以便在模型执行时为其分配正确的内存空间。如果输出形状不固定，那么 ONNX 运行时就需要在运行时动态调整输出形状，这将使得模型在部署时的性能受到影响。

因此，在转换 PyTorch 模型为 ONNX 模型时，你需要手动指定每个输出的固定形状，以便在执行时能够顺利运行。