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

import torch
import torch.nn as nn
from torchvision import models

class 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)

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"}}
)