如何解决PyTorch ResNet50模型导出ONNX时动态Batch Size的问题？

wufei123 2025-03-12 阅读:168 评论:0

本文探讨如何将PyTorch ResNet50模型导出为ONNX格式，并使其支持动态Batch Size。许多深度学习模型需要导出为ONNX格式以便在不同框架中部署，而处理动态Batch Size是其中一个常见挑战。在将自定义Res...

如何解决pytorch resnet50模型导出onnx时动态batch size的问题？

本文探讨如何将PyTorch ResNet50模型导出为ONNX格式，并使其支持动态Batch Size。许多深度学习模型需要导出为ONNX格式以便在不同框架中部署，而处理动态Batch Size是其中一个常见挑战。

在将自定义ResNet50模型（包含GEM和L2N自定义层）导出为ONNX时，原始代码因torch.onnx.export函数失败而报错，错误信息指向prim::listconstruct节点的非常量输入，这与GEM层中动态计算的p值以及模型中未使用的属性有关。

为了解决这个问题，我们需要进行以下修改：

首先，移除模型中未使用的属性，简化模型结构，避免对ONNX形状推断的干扰。

其次，修改GEM层，将self.p从参数类型改为直接赋值常量值。参数类型在ONNX导出过程中可能导致动态形状问题。修改后的GEM层代码如下：

PHP

class gem(nn.Module):
    def __init__(self, p=3, eps=1e-6):
        super(gem, self).__init__()
        self.p = p  # 直接赋值常量
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return gem_op(x, p=self.p, eps=self.eps) # 假设gem_op是一个实现了GEM运算的函数

最后，修改后的ResNet50模型代码（示例）：

PHP

class ImageRetrievalNet(nn.Module):
    def __init__(self, dim: int = 512):
        super(ImageRetrievalNet, self).__init__()
        # ... ResNet50 模型构建代码 ...
        self.pool = GeM()
        self.whiten = nn.Linear(2048, dim, bias=True)
        self.norm = L2N()

    def forward(self, x: Tensor):
        # ... 模型前向传播代码 ...
        return o.permute(1, 0)

通过这些修改，GEM层的p值成为静态常量，避免了动态形状问题。移除未使用的属性简化了模型，使ONNX导出过程更加顺利。使用torch.onnx.export函数导出时，建议设置opset_version为12或更高版本，以确保更好的兼容性。这样就能成功导出支持动态Batch Size的ONNX模型。

以上就是如何解决PyTorch ResNet50模型导出ONNX时动态Batch Size的问题？的详细内容，更多请关注知识资源分享宝库其它相关文章！