How to convert PyTorch Mask R-CNN model to Core ML

Question

I'm looking for a way to convert my custom Mask R-CNN model to Core ML.

I'm facing a problem when trying to trace the model.

Here is an example of a basic model that works fine if I try to convert just backbone models.resnet50(pretrained=True) instead of full model models.detection.maskrcnn_resnet50_fpn_v2(pretrained=True)

import torch
import torchvision.models as models
import coremltools as ct

# Load the Mask R-CNN model
model = models.detection.maskrcnn_resnet50_fpn_v2(pretrained=True)
model.eval()

# Trace the model with a dummy input
dummy_input = torch.randn(1, 3, 800, 800)
traced_model = torch.jit.trace(model, dummy_input)

# Convert the traced model to Core ML and assign custom names
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input", shape=dummy_input.shape)],
    outputs=[ct.TensorType(name="output")]
)

# Add descriptions
mlmodel.input_description['input'] = "Input image"
mlmodel.output_description['output'] = "Detected features"

# Save the model
mlmodel.save("MaskRCNN.mlpackage")

For this script I got an error RuntimeError: Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions

Also I have tried to script the model instead of tracing it, but in that case, I have problems with the Core ML converter itself as it experimentally supports script models.

import torch
import torchvision.models as models
import coremltools as ct

# Load the Mask R-CNN model
model = models.detection.maskrcnn_resnet50_fpn_v2(pretrained=True)
model.eval()

scripted_model = torch.jit.script(model)
scripted_model.eval()

# Trace the model with a dummy input
dummy_input = torch.randn(1, 3, 800, 800)

# Convert the scripted model to Core ML and assign custom names
mlmodel = ct.convert(
    scripted_model,
    inputs=[ct.TensorType(name="input", shape=dummy_input.shape)],
    outputs=[ct.TensorType(name="output")]
)

# Add descriptions
mlmodel.input_description['input'] = "Input image"
mlmodel.output_description['output'] = "Detected features"

# Save the model
mlmodel.save("MaskRCNN.mlpackage")

This leads to the error KeyError: images.7 defined in (%images.7 : __torch__.torchvision.models.detection.image_list.ImageList, %targets.31 : Dict(str, Tensor)[]? = prim::TupleUnpack(%413) )

Is there any way to convert the Mask R-CNN model to Core ML?