Create model/modeling_maskrcnn.py

model/modeling_maskrcnn.py

@@ -0,0 +1,42 @@
+import torch
+from torchvision.models.detection import maskrcnn_resnet50_fpn
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
+
+from transformers import PreTrainedModel
+from .configuration_maskrcnn import MaskRCNNConfig
+
+
+class MaskRCNNForInstanceSegmentation(PreTrainedModel):
+    config_class = MaskRCNNConfig
+
+    def __init__(self, config: MaskRCNNConfig):
+        super().__init__(config)  
+        self.config = config
+
+        model = maskrcnn_resnet50_fpn(weights=None)
+
+        # box head
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(
+            in_features, config.num_classes
+        )
+
+        # mask head
+        in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
+        model.roi_heads.mask_predictor = MaskRCNNPredictor(
+            in_features_mask,
+            config.hidden_layer,
+            config.num_classes
+        )
+
+        self.model = model
+
+    def forward(self, images, targets=None):
+        """
+        Train:
+            returns dict(loss_classifier, loss_box_reg, loss_mask, ...)
+        Eval:
+            returns List[Dict(boxes, labels, scores, masks)]
+        """
+        return self.model(images, targets)