本文探讨如何将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层代码如下：
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模型代码（示例）：
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模型。
