PyTorch ResNet50模型导出ONNX时如何解决动态batch_size难题？-Python教程-PHP中文网

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

DDD

发布： 2025-03-20 09:40:02

原创

322人浏览过

解决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参数改为直接赋值的常量，不再作为可学习参数：

百灵大模型

蚂蚁集团自研的多模态AI大模型系列

177

查看详情

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导出过程中直接使用类名调用。

以上就是PyTorch ResNet50模型导出ONNX时如何解决动态batch_size难题？的详细内容，更多请关注php中文网其它相关文章！