MaskRCNN-ConvNeXt Instance Segmentation (#2292)

* add swin-t

* update config

* simplify data pipeline and use SGD optimizer

* update model

* add convnext-t

* update convnext input size

* update custom mask fcn head

* add x101

* update

* remove swin-t and ResNeXt template

* style change

* rename convnext

* update compressed model yaml

* modify nncf compression yaml

* update doc

* update structure

* update pretrained model path

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/data_pipeline.py

Co-authored-by: Sungman Cho <[email protected]>

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/tile_pipeline.py

Co-authored-by: Sungman Cho <[email protected]>

* change Gflops

* update doc

* fix model url path

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/tile_pipeline.py

Co-authored-by: Jaeguk Hyun <[email protected]>

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/model.py

Co-authored-by: Jaeguk Hyun <[email protected]>

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/deployment_tile_classifier.py

Co-authored-by: Jaeguk Hyun <[email protected]>

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/data_pipeline.py

Co-authored-by: Jaeguk Hyun <[email protected]>

* Update src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/__init__.py

Co-authored-by: Jaeguk Hyun <[email protected]>

* fix according to review

* update hpo

* include convnext tests

---------

Co-authored-by: Sungman Cho <[email protected]>
Co-authored-by: Jaeguk Hyun <[email protected]>

Author: 3 people

docs/source/guide/explanation/algorithms/segmentation/instance_segmentation.rst

@@ -58,15 +58,21 @@ Models
 
 We support the following ready-to-use model templates:
 
-+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
-| Template ID                                                                                                                                                                                                                                    | Name                       | Complexity (GFLOPs) | Model size (MB) |
-+================================================================================================================================================================================================================================================+============================+=====================+=================+
-| `Custom_Counting_Instance_Segmentation_MaskRCNN_EfficientNetB2B <https://github.com/openvinotoolkit/training_extensions/blob/develop/src/otx/algorithms/detection/configs/instance_segmentation/efficientnetb2b_maskrcnn/template.yaml>`_      | MaskRCNN-EfficientNetB2B   | 68.48               | 13.27           |
-+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
-| `Custom_Counting_Instance_Segmentation_MaskRCNN_ResNet50 <https://github.com/openvinotoolkit/training_extensions/blob/develop/src/otx/algorithms/detection/configs/instance_segmentation/resnet50_maskrcnn/template.yaml>`_                    | MaskRCNN-ResNet50          | 533.80              | 177.90          |
-+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
++--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
+| Template ID                                                                                                                                                                                                                                | Name                       | Complexity (GFLOPs) | Model size (MB) |
++============================================================================================================================================================================================================================================+============================+=====================+=================+
+| `Custom_Counting_Instance_Segmentation_MaskRCNN_EfficientNetB2B <https://github.com/openvinotoolkit/training_extensions/blob/develop/src/otx/algorithms/detection/configs/instance_segmentation/efficientnetb2b_maskrcnn/template.yaml>`_      | MaskRCNN-EfficientNetB2B   | 68.48           | 13.27           |
++--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
+| `Custom_Counting_Instance_Segmentation_MaskRCNN_ResNet50 <https://github.com/openvinotoolkit/training_extensions/blob/develop/src/otx/algorithms/detection/configs/instance_segmentation/resnet50_maskrcnn/template.yaml>`_                    | MaskRCNN-ResNet50          | 533.80          | 177.90          |
++--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
+| `Custom_Counting_Instance_Segmentation_MaskRCNN_ConvNeXt <https://github.com/openvinotoolkit/training_extensions/blob/develop/src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/template.yaml>`_                    | MaskRCNN-ConvNeXt          | 266.78          | 192.4          |
++--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------+---------------------+-----------------+
 
-``MaskRCNN-ResNet50`` uses `ResNet-50 <https://arxiv.org/abs/1512.03385>`_ as the backbone network for the image features extraction. It has more parameters and FLOPs and needs more time to train, meanwhile providing superior performance in terms of accuracy. ``MaskRCNN-EfficientNetB2B`` uses `EfficientNet-B2 <https://arxiv.org/abs/1905.11946>`_ as the backbone network. It is a good trade-off between accuracy and speed. It is a better choice when training time and computational cost are in priority.
+MaskRCNN-ResNet50 utilizes the `ResNet-50 <https://arxiv.org/abs/1512.03385>`_ architecture as the backbone network for extracting image features. This choice of backbone network results in a higher number of parameters and FLOPs, which consequently requires more training time. However, the model offers superior performance in terms of accuracy.
+
+On the other hand, MaskRCNN-EfficientNetB2B employs the `EfficientNet-B2 <https://arxiv.org/abs/1905.11946>`_ architecture as the backbone network. This selection strikes a balance between accuracy and speed, making it a preferable option when prioritizing training time and computational cost.
+
+Recently, we have made updates to MaskRCNN-ConvNeXt, incorporating the `ConvNeXt backbone <https://arxiv.org/abs/2201.03545>`_. Through our experiments, we have observed that this variant achieves better accuracy compared to MaskRCNN-ResNet50 while utilizing less GPU memory. However, it is important to note that the training time and inference duration may slightly increase. If minimizing training time is a significant concern, we recommend considering a switch to MaskRCNN-EfficientNetB2B.
 
 .. In the table below the `mAP <https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient>`_ metric on some academic datasets using our :ref:`supervised pipeline <instance_segmentation_supervised_pipeline>` is presented. The results were obtained on our templates without any changes. We use 1024x1024 image resolution, for other hyperparameters, please, refer to the related template. We trained each model with single Nvidia GeForce RTX3090.
 
@@ -77,6 +83,8 @@ We support the following ready-to-use model templates:
 .. +---------------------------+--------------+------------+-----------------+
 .. | MaskRCNN-ResNet50         | N/A          | N/A        | N/A             |
 .. +---------------------------+--------------+------------+-----------------+
+.. | MaskRCNN-ConvNeXt         | N/A          | N/A        | N/A             |
+.. +---------------------------+--------------+------------+-----------------+
 
 .. *******************
 .. Tiling Pipeline

docs/source/guide/tutorials/base/how_to_train/instance_segmentation.rst

@@ -136,6 +136,7 @@ The list of supported templates for instance segmentation is available with the
   +-----------------------+----------------------------------------------------------------+--------------------------+---------------------------------------------------------------------------------------------------+
   | INSTANCE_SEGMENTATION |    Custom_Counting_Instance_Segmentation_MaskRCNN_ResNet50     |    MaskRCNN-ResNet50     |     src/otx/algorithms/detection/configs/instance_segmentation/resnet50_maskrcnn/template.yaml    |
   | INSTANCE_SEGMENTATION | Custom_Counting_Instance_Segmentation_MaskRCNN_EfficientNetB2B | MaskRCNN-EfficientNetB2B | src/otx/algorithms/detection/configs/instance_segmentation/efficientnetb2b_maskrcnn/template.yaml |
+  | INSTANCE_SEGMENTATION | Custom_Counting_Instance_Segmentation_MaskRCNN_ConvNeXt        | MaskRCNN-ConvNeXt        | src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/template.yaml        |
   +-----------------------+----------------------------------------------------------------+--------------------------+---------------------------------------------------------------------------------------------------+
 
 2. We need to create

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/__init__.py

@@ -0,0 +1,4 @@
+"""Initialization of ConvNeXt-T-MaskRCNN model for Instance-Segmentation Task."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/compression_config.json

@@ -0,0 +1,41 @@
+{
+  "base": {
+    "find_unused_parameters": true,
+    "nncf_config": {
+      "target_metric_name": "mAP",
+      "input_info": {
+        "sample_size": [1, 3, 1024, 1024]
+      },
+      "compression": [],
+      "log_dir": "/tmp"
+    }
+  },
+  "nncf_quantization": {
+    "optimizer": {
+      "lr": 0.0005
+    },
+    "nncf_config": {
+      "compression": [
+        {
+          "algorithm": "quantization",
+          "initializer": {
+            "range": {
+              "num_init_samples": 1000
+            },
+            "batchnorm_adaptation": {
+              "num_bn_adaptation_samples": 1000
+            }
+          }
+        }
+      ],
+      "accuracy_aware_training": {
+        "mode": "early_exit",
+        "params": {
+          "maximal_absolute_accuracy_degradation": 0.01,
+          "maximal_total_epochs": 20
+        }
+      }
+    }
+  },
+  "order_of_parts": ["nncf_quantization"]
+}

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/data_pipeline.py

@@ -0,0 +1,64 @@
+"""Data Pipeline of ConvNeXt model for Instance-Seg Task."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# pylint: disable=invalid-name
+
+__img_size = (1024, 1024)
+
+# TODO: A comparison experiment is needed to determine which value is appropriate for to_rgb.
+__img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+
+train_pipeline = [
+    dict(type="LoadImageFromFile"),
+    dict(type="LoadAnnotations", with_bbox=True, with_mask=True, poly2mask=False),
+    dict(type="Resize", img_scale=__img_size, keep_ratio=False),
+    dict(type="RandomFlip", flip_ratio=0.5),
+    dict(type="Normalize", **__img_norm_cfg),
+    dict(type="DefaultFormatBundle"),
+    dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels", "gt_masks"]),
+]
+
+test_pipeline = [
+    dict(type="LoadImageFromFile"),
+    dict(
+        type="MultiScaleFlipAug",
+        img_scale=__img_size,
+        flip=False,
+        transforms=[
+            dict(type="Resize", keep_ratio=False),
+            dict(type="RandomFlip"),
+            dict(type="Normalize", **__img_norm_cfg),
+            dict(type="ImageToTensor", keys=["img"]),
+            dict(type="Collect", keys=["img"]),
+        ],
+    ),
+]
+
+__dataset_type = "CocoDataset"
+
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=__dataset_type,
+        ann_file="data/coco/annotations/instances_train2017.json",
+        img_prefix="data/coco/train2017",
+        pipeline=train_pipeline,
+    ),
+    val=dict(
+        type=__dataset_type,
+        test_mode=True,
+        ann_file="data/coco/annotations/instances_val2017.json",
+        img_prefix="data/coco/val2017",
+        pipeline=test_pipeline,
+    ),
+    test=dict(
+        type=__dataset_type,
+        test_mode=True,
+        ann_file="data/coco/annotations/instances_val2017.json",
+        img_prefix="data/coco/val2017",
+        pipeline=test_pipeline,
+    ),
+)

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/deployment.py

@@ -0,0 +1,9 @@
+"""MMDployment config of Resnet model for Instance-Seg Task."""
+
+_base_ = ["../../base/deployments/base_instance_segmentation_dynamic.py"]
+
+scale_ir_input = True
+
+ir_config = dict(
+    output_names=["boxes", "labels", "masks"],
+)

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/deployment_tile_classifier.py

@@ -0,0 +1,22 @@
+"""MMDeploy config partitioning ConvNeXt-T MaskRCNN model to tile classifier and MaskRCNN model."""
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+_base_ = ["./deployment.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels", "masks", "tile_prob"],
+)
+
+partition_config = dict(
+    type="tile_classifier",
+    apply_marks=True,
+    partition_cfg=[
+        dict(
+            save_file="tile_classifier.onnx",
+            start=["tile_classifier:input"],
+            end=["tile_classifier:output"],
+            output_names=["tile_prob"],
+        )
+    ],
+)

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/hpo_config.yaml

@@ -0,0 +1,16 @@
+metric: mAP
+search_algorithm: asha
+early_stop: None
+hp_space:
+  learning_parameters.learning_rate:
+    param_type: qloguniform
+    range:
+      - 0.0001
+      - 0.01
+      - 0.0001
+  learning_parameters.batch_size:
+    param_type: qloguniform
+    range:
+      - 2
+      - 6
+      - 2

src/otx/algorithms/detection/configs/instance_segmentation/convnext_maskrcnn/model.py

@@ -0,0 +1,133 @@
+"""Model configuration of ConvNeXt-T-MaskRCNN model for Instance-Seg Task."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# pylint: disable=invalid-name
+
+_base_ = [
+    "../../../../../recipes/stages/instance-segmentation/incremental.py",
+    "../../base/models/detector.py",
+]
+
+task = "instance-segmentation"
+
+model = dict(
+    type="CustomMaskRCNN",
+    backbone=dict(
+        type="mmcls.ConvNeXt",
+        arch="tiny",
+        out_indices=[0, 1, 2, 3],
+        drop_path_rate=0.4,
+        layer_scale_init_value=1.0,
+        gap_before_final_norm=False,
+    ),
+    neck=dict(type="FPN", in_channels=[96, 192, 384, 768], out_channels=256, num_outs=5),
+    rpn_head=dict(
+        type="RPNHead",
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(type="AnchorGenerator", scales=[8], ratios=[0.5, 1.0, 2.0], strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(type="DeltaXYWHBBoxCoder", target_means=[0.0, 0.0, 0.0, 0.0], target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(type="CrossEntropyLoss", use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type="L1Loss", loss_weight=1.0),
+    ),
+    roi_head=dict(
+        type="CustomRoIHead",
+        bbox_roi_extractor=dict(
+            type="SingleRoIExtractor",
+            roi_layer=dict(type="RoIAlign", output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32],
+        ),
+        bbox_head=dict(
+            type="Shared2FCBBoxHead",
+            in_channels=256,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type="DeltaXYWHBBoxCoder", target_means=[0.0, 0.0, 0.0, 0.0], target_stds=[0.1, 0.1, 0.2, 0.2]
+            ),
+            reg_class_agnostic=False,
+            loss_cls=dict(type="CrossEntropyLoss", use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type="L1Loss", loss_weight=1.0),
+        ),
+        mask_roi_extractor=dict(
+            type="SingleRoIExtractor",
+            roi_layer=dict(type="RoIAlign", output_size=14, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32],
+        ),
+        mask_head=dict(
+            type="CustomFCNMaskHead",
+            num_convs=4,
+            in_channels=256,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(type="CrossEntropyLoss", use_mask=True, loss_weight=1.0),
+        ),
+    ),
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type="CustomMaxIoUAssigner",
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1,
+                gpu_assign_thr=300,
+            ),
+            sampler=dict(type="RandomSampler", num=256, pos_fraction=0.5, neg_pos_ub=-1, add_gt_as_proposals=False),
+            allowed_border=-1,
+            pos_weight=-1,
+            debug=False,
+        ),
+        rpn_proposal=dict(
+            nms_across_levels=False,
+            nms_pre=2000,
+            max_per_img=1000,
+            nms=dict(type="nms", iou_threshold=0.7),
+            min_bbox_size=0,
+        ),
+        rcnn=dict(
+            assigner=dict(
+                type="CustomMaxIoUAssigner",
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=True,
+                ignore_iof_thr=-1,
+                gpu_assign_thr=300,
+            ),
+            sampler=dict(type="RandomSampler", num=512, pos_fraction=0.25, neg_pos_ub=-1, add_gt_as_proposals=True),
+            mask_size=28,
+            pos_weight=-1,
+            debug=False,
+        ),
+    ),
+    test_cfg=dict(
+        rpn=dict(
+            nms_across_levels=False,
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type="nms", iou_threshold=0.7),
+            min_bbox_size=0,
+        ),
+        rcnn=dict(
+            score_thr=0.05, nms=dict(type="nms", iou_threshold=0.5, max_num=100), max_per_img=100, mask_thr_binary=0.5
+        ),
+    ),
+)
+
+load_from = "https://storage.openvinotoolkit.org/\
+repositories/openvino_training_extensions/\
+models/instance_segmentation/\
+mask_rcnn_convnext-t_p4_w7_fpn_fp16.pth"
+
+evaluation = dict(interval=1, metric="mAP", save_best="mAP", iou_thr=[0.5])
+ignore = True
+
+custom_imports = dict(imports=["mmcls.models"], allow_failed_imports=False)
+fp16 = dict(loss_scale=dict(init_scale=512.0))