如何将PyTorch ResNet50模型导出为支持动态batch size的ONNX格式？

本文阐述如何将PyTorch ResNet50模型导出为支持动态batch size的ONNX格式。 我们将解决导出过程中可能遇到的与动态batch size相关的难题。 原始代码中存在一些问题，例如imageretrievalnet类中未使用的self.lwhiten属性和gem类中使用parameter类型的self.p属性，这些都可能导致ONNX导出失败。

ONNX导出对动态形状和特定PyTorch操作符的处理有限制。gem类中self.p作为可学习参数，以及imageretrievalnet类中冗余属性，会增加导出复杂度，导致形状推断错误。

为了解决这些问题，我们需要修改代码：移除imageretrievalnet类中无用的self.lwhiten属性，并将gem类中的self.p属性改为常量赋值。 修改后的代码如下：

import torchimport torch.nn as nnfrom torchvision import modelsclass gem(nn.Module):    def __init__(self, p=3, eps=1e-6):        super(gem, self).__init__()        self.p = p  # p设置为常量        self.eps = eps    def forward(self, x: torch.Tensor) -> torch.Tensor:        return torch.pow(torch.mean(torch.pow(x, self.p), dim=1, keepdim=True) + self.eps, 1/self.p)class l2n(nn.Module):    def __init__(self):        super(l2n, self).__init__()    def forward(self, x: torch.Tensor) -> torch.Tensor:        return x / torch.norm(x, p=2, dim=1, keepdim=True)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)        self.norm = l2n()    def forward(self, x: torch.Tensor) -> 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)

登录后复制