PyTorchResNet50模型导出ONNX时如何解决动态batch
解决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 = 4input_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导出过程中直接使用类名调用。
以上就是PyTorch ResNet50模型导出ONNX时如何解决动态batch_size难题？的详细内容，更多请关注范的资源库其它相关文章！
