Customize pytorch model export to ONNX

vision
1

I am trying to export pretrained Mask R-CNN model to ONNX format. Since this model in basic configuration has following structure (here I added batch_size as dynamic axes):

enter image description here

I want to customize my model and add batch_size to output (it means I need to add new dim to each of the outputs).

I wrote following code to make it possible:

class MaskRCNNModel(torch.nn.Module):
  def __init__(self):
    super(MaskRCNNModel, self).__init__()
    self.model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights='DEFAULT')
    in_features = self.model.roi_heads.box_predictor.cls_score.in_features
    self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=7)
    self.model.load_state_dict(torch.load("saved_dict.torch"))

  def forward(self, input):
    outputs = self.model.forward(input)
    boxes = []
    labels = []
    scores = []
    masks = []
    for result in outputs:
        box, label, score, mask = result.values()
        boxes.append(box)
        labels.append(label)
        scores.append(score)
        masks.append(mask)
    
    return boxes, labels, scores, masks

maskrcnn_model = MaskRCNNModel()
maskrcnn_model.eval()
maskrcnn_model.to(device)

x = torch.rand(1, 3, 512, 512)
x = x.to(device)

maskrcnn_model(x)

torch.onnx.export(maskrcnn_model,
                  x,
                  "base_model_100_epochs.onnx",
                  opset_version=11,
                  input_names=["input"],
                  output_names=["boxes", "labels", "scores", "masks"])

but the code above doesn’t change any export parameters. The structure of output stays the same:

enter image description here

What should I do to customize forward method to be able to add batch_size into ONNX model output?