PyTorch ResNet50模型导出ONNX时如何解决动态batch

解决pytorch resnet50模型导出onnx时动态batch_size难题

本文介绍如何将基于ResNet50的PyTorch模型导出为ONNX格式，重点解决动态batch_size导致的导出问题。原始代码中，imageretrievalnet类和gem类存在一些与ONNX导出不兼容的因素，主要包括gem类中可学习参数self.p以及imageretrievalnet类中未使用的self.lwhiten属性。这些动态元素阻碍了ONNX的shape推断，导致导出失败。

PyTorch ResNet50模型导出ONNX时如何解决动态batch_size难题？

为了解决这个问题，我们需要修改这两个类以适应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导出过程中直接使用类名调用。

文章版权归作者所有，未经允许请勿转载。

THE END