本文阐述如何将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)
通过这些修改，我们解决了ONNX导出过程中遇到的动态形状和参数类型不匹配问题。 现在，可以使用以下代码将修改后的模型导出为支持动态batch size的ONNX格式：
model = ImageRetrievalNet()batch_size = 1 # 使用1进行测试，实际部署时batch size可变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"}})
修改后的代码能够成功导出ONNX模型并支持动态batch size，主要是因为消除了阻碍动态形状推断的因素，从而使ONNX导出过程顺利完成。 请确保已安装必要的库：torch, torchvision, onnx.
以上就是如何将PyTorch ResNet50模型导出为支持动态batch size的ONNX格式？的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1358960.html
微信扫一扫 
支付宝扫一扫 
