export-mask-rcnn-onnx.py

import torch, torchvision
from onnxruntime import InferenceSession

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
x = torch.rand(1, 3, 360, 640)

torch.onnx.export(
    model,
    x,  # ONNX requires fixed input size
    "test-mask-rcnn-export.onnx",
    do_constant_folding=True,
    dynamic_axes={
        "images_tensors": [0, 1, 2, 3],
        "boxes": [0, 1],
        "labels": [0],
        "scores": [0],
        "masks": [0, 1, 2, 3],
    },
    # dynamic_axes={"input_image": {0: "sequence"}, "output": {0: "sequence"}},
    input_names=["input_image"],
    output_names=["bbox_coords", "bbox_labels", "bbox_scores", "bbox_masks"],
    opset_version=11,  # opset_version 11 required for Mask R-CNN
)

session = InferenceSession("test-mask-rcnn-export.onnx")
session.run(None, {"input_image": x.detach().cpu().numpy()} )