解决pytorch resnet50模型导出onnx时动态batch_size难题
本文介绍如何将基于ResNet50的PyTorch模型导出为ONNX格式，重点解决动态batch_size导致的导出问题。原始代码中，imageretrievalnet类和gem类存在一些与ONNX导出不兼容的因素，主要包括gem类中可学习参数self.p以及imageretrievalnet类中未使用的self.lwhiten属性。这些动态元素阻碍了ONNX的shape推断，导致导出失败。
为了解决这个问题，我们需要修改这两个类以适应ONNX导出流程。具体修改如下：
首先，修改gem类，将self.p参数改为直接赋值的常量，不再作为可学习参数：
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函数，避免直接使用类名调用
然后，简化imageretrievalnet类，去除未使用的self.lwhiten属性：
class imageretrievalnet(nn.Module):
def __init__(self, dim: int = 512):
super(imageretrievalnet, self).__init__()
resnet50_model = models.resnet50()
features = list(resnet50_model.children())[:-2]
self.features = nn.Sequential(*features)
self.pool = gem()
self.whiten = nn.Linear(2048, dim, bias=True) # 使用nn.Linear
self.norm = l2n()
def forward(self, x: torch.Tensor):
o: torch.Tensor = self.features(x)
pooled_t = self.pool(o)
normed_t: torch.Tensor = self.norm(pooled_t)
o: torch.Tensor = normed_t.squeeze(-1).squeeze(-1)
if self.whiten is not None:
whitened_t = self.whiten(o)
normed_t: torch.Tensor = self.norm(whitened_t)
o = normed_t
return o.permute(1, 0)
通过以上修改，消除了动态参数带来的不确定性，使ONNX导出能够顺利进行。 使用修改后的imageretrievalnet类，并利用torch.onnx.export函数，指定dynamic_axes参数处理动态batch_size，即可成功导出ONNX模型：
model = imageretrievalnet()
batch_size = 4
input_shape = (batch_size, 3, 224, 224)
input_data = torch.randn(input_shape)
torch.onnx.export(
model,
input_data,
"resnet50.onnx",
input_names=["input"], output_names=["output"],
opset_version=12,
dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)
记住根据实际情况调整opset_version参数。 通过这些修改，即可成功导出支持动态batch_size的ResNet50 ONNX模型。 请注意，代码中添加了gem_op函数的假设，该函数应该实现gem类的功能，以避免在ONNX导出过程中直接使用类名调用。
