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

  • Z时代
  • 2024-02-13
  • 分类：IT

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

最近在把 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 显存大小等等做出调整的，而不是固定死的才对

回答：

                  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" 只是为了可以方便理解模型。

回答：

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

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

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

以上是 为什么 pytorch 转成 onnx 的时候，需要固定输出形状？ 的全部内容， 来源链接： utcz.com/p/938771.html

回到顶部