add test units for train SAN

tests/test_models/test_assigners/test_hungarian_assigner.py

@@ -0,0 +1,77 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from unittest import TestCase
+
+import torch
+from mmengine.structures import InstanceData
+
+from mmseg.models.assigners import HungarianAssigner
+
+
+class TestHungarianAssigner(TestCase):
+
+    def test_init(self):
+        with self.assertRaises(AssertionError):
+            HungarianAssigner([])
+
+    def test_hungarian_match_assigner(self):
+        assigner = HungarianAssigner([
+            dict(type='ClassificationCost', weight=2.0),
+            dict(type='CrossEntropyLossCost', weight=5.0, use_sigmoid=True),
+            dict(type='DiceCost', weight=5.0, pred_act=True, eps=1.0)
+        ])
+        num_classes = 3
+        num_masks = 10
+        num_points = 20
+        gt_instances = InstanceData()
+        gt_instances.labels = torch.randint(0, num_classes, (num_classes, ))
+        gt_instances.masks = torch.randint(0, 2, (num_classes, num_points))
+        pred_instances = InstanceData()
+        pred_instances.scores = torch.rand((num_masks, num_classes))
+        pred_instances.masks = torch.rand((num_masks, num_points))
+
+        matched_quiery_inds, matched_label_inds = \
+            assigner.assign(pred_instances, gt_instances)
+        unique_quiery_inds = torch.unique(matched_quiery_inds)
+        unique_label_inds = torch.unique(matched_label_inds)
+        self.assertTrue(len(unique_quiery_inds) == len(matched_quiery_inds))
+        self.assertTrue(
+            torch.equal(unique_label_inds, torch.arange(0, num_classes)))
+
+    def test_cls_match_cost(self):
+        num_classes = 3
+        num_masks = 10
+        gt_instances = InstanceData()
+        gt_instances.labels = torch.randint(0, num_classes, (num_classes, ))
+        pred_instances = InstanceData()
+        pred_instances.scores = torch.rand((num_masks, num_classes))
+
+        # test ClassificationCost
+        assigner = HungarianAssigner(dict(type='ClassificationCost'))
+        matched_quiery_inds, matched_label_inds = \
+            assigner.assign(pred_instances, gt_instances)
+        unique_quiery_inds = torch.unique(matched_quiery_inds)
+        unique_label_inds = torch.unique(matched_label_inds)
+        self.assertTrue(len(unique_quiery_inds) == len(matched_quiery_inds))
+        self.assertTrue(
+            torch.equal(unique_label_inds, torch.arange(0, num_classes)))
+
+    def test_mask_match_cost(self):
+        num_classes = 3
+        num_masks = 10
+        num_points = 20
+        gt_instances = InstanceData()
+        gt_instances.masks = torch.randint(0, 2, (num_classes, num_points))
+        pred_instances = InstanceData()
+        pred_instances.masks = torch.rand((num_masks, num_points))
+
+        # test DiceCost
+        assigner = HungarianAssigner(
+            dict(type='DiceCost', pred_act=True, eps=1.0))
+        assign_result = assigner.assign(pred_instances, gt_instances)
+        self.assertTrue(len(assign_result[0]) == len(assign_result[1]))
+
+        # test CrossEntropyLossCost
+        assigner = HungarianAssigner(
+            dict(type='CrossEntropyLossCost', use_sigmoid=True))
+        assign_result = assigner.assign(pred_instances, gt_instances)
+        self.assertTrue(len(assign_result[0]) == len(assign_result[1]))

tests/test_models/test_heads/test_san_head.py

@@ -5,60 +5,95 @@
 
 from mmseg.models.decode_heads import SideAdapterCLIPHead
 from mmseg.structures import SegDataSample
-from mmseg.utils import SampleList
 from .utils import list_to_cuda
 
 
 def test_san_head():
     H, W = (64, 64)
-    clip_channels = 768
+    clip_channels = 64
     img_channels = 4
-    num_queries = 100
-    out_dims = 512
+    num_queries = 40
+    out_dims = 64
     num_classes = 19
     cfg = dict(
         num_classes=num_classes,
+        deep_supervision_idxs=[4],
         san_cfg=dict(
             in_channels=img_channels,
-            embed_dims=240,
+            embed_dims=128,
             clip_channels=clip_channels,
             num_queries=num_queries,
-            cfg_encoder=dict(num_encode_layer=12, mlp_ratio=4, num_heads=12),
+            cfg_encoder=dict(num_encode_layer=4, mlp_ratio=2, num_heads=2),
             cfg_decoder=dict(
-                num_heads=12,
+                num_heads=4,
                 num_layers=1,
-                embed_channels=256,
-                mlp_channels=256,
-                num_mlp=3,
+                embed_channels=32,
+                mlp_channels=32,
+                num_mlp=2,
                 rescale=True)),
         maskgen_cfg=dict(
             sos_token_num=num_queries,
             embed_dims=clip_channels,
-            out_dims=out_dims))
-    cfg = Config(cfg)
+            out_dims=out_dims,
+            num_heads=4,
+            mlp_ratio=2),
+        train_cfg=dict(
+            num_points=100,
+            oversample_ratio=3.0,
+            importance_sample_ratio=0.75,
+            assigner=dict(
+                type='HungarianAssigner',
+                match_costs=[
+                    dict(type='ClassificationCost', weight=2.0),
+                    dict(
+                        type='CrossEntropyLossCost',
+                        weight=5.0,
+                        use_sigmoid=True),
+                    dict(type='DiceCost', weight=5.0, pred_act=True, eps=1.0)
+                ])),
+        loss_decode=[
+            dict(
+                type='CrossEntropyLoss',
+                loss_name='loss_cls_ce',
+                loss_weight=2.0,
+                class_weight=[1.0] * num_classes + [0.1]),
+            dict(
+                type='CrossEntropyLoss',
+                use_sigmoid=True,
+                loss_name='loss_mask_ce',
+                loss_weight=5.0),
+            dict(
+                type='DiceLoss',
+                ignore_index=None,
+                naive_dice=True,
+                eps=1,
+                loss_name='loss_mask_dice',
+                loss_weight=5.0)
+        ])
 
+    cfg = Config(cfg)
     head = SideAdapterCLIPHead(**cfg)
 
-    inputs = torch.rand((1, img_channels, H, W))
+    inputs = torch.rand((2, img_channels, H, W))
     clip_feature = [[
-        torch.rand((1, clip_channels, H // 2, W // 2)),
-        torch.rand((1, clip_channels))
+        torch.rand((2, clip_channels, H // 2, W // 2)),
+        torch.rand((2, clip_channels))
     ],
                     [
-                        torch.rand((1, clip_channels, H // 2, W // 2)),
-                        torch.rand((1, clip_channels))
+                        torch.rand((2, clip_channels, H // 2, W // 2)),
+                        torch.rand((2, clip_channels))
                     ],
                     [
-                        torch.rand((1, clip_channels, H // 2, W // 2)),
-                        torch.rand((1, clip_channels))
+                        torch.rand((2, clip_channels, H // 2, W // 2)),
+                        torch.rand((2, clip_channels))
                     ],
                     [
-                        torch.rand((1, clip_channels, H // 2, W // 2)),
-                        torch.rand((1, clip_channels))
+                        torch.rand((2, clip_channels, H // 2, W // 2)),
+                        torch.rand((2, clip_channels))
                     ]]
     class_embed = torch.rand((num_classes + 1, out_dims))
 
-    data_samples: SampleList = []
+    data_samples = []
     for i in range(2):
         data_sample = SegDataSample()
         img_meta = {}
@@ -68,23 +103,24 @@ def test_san_head():
             data=torch.randint(0, num_classes, (1, H, W)))
         data_sample.set_metainfo(img_meta)
         data_samples.append(data_sample)
+
     batch_img_metas = []
     for data_sample in data_samples:
         batch_img_metas.append(data_sample.metainfo)
 
     if torch.cuda.is_available():
         head = head.cuda()
         data = list_to_cuda([inputs, clip_feature, class_embed])
+        for data_sample in data_samples:
+            data_sample.gt_sem_seg.data = data_sample.gt_sem_seg.data.cuda()
     else:
         data = [inputs, clip_feature, class_embed]
 
-    with torch.no_grad():
-        mask_props, mask_logits = head.forward(data)
-    assert len(mask_props) == len(mask_logits)
-    for mask_prop, mask_logit in zip(mask_props, mask_logits):
-        assert mask_prop.shape == (1, num_queries, H // 16, W // 16)
-        assert mask_logit.shape == (1, num_queries, num_classes + 1)
+    # loss test
+    loss_dict = head.loss(data, data_samples, None)
+    assert isinstance(loss_dict, dict)
 
+    # prediction test
     with torch.no_grad():
         seg_logits = head.predict(data, batch_img_metas, None)
-    assert seg_logits.shape == torch.Size((1, num_classes, H, W))
+    assert seg_logits.shape == torch.Size((2, num_classes, H, W))

tests/test_models/test_segmentors/test_multimodal_encoder_decoder.py

@@ -2,7 +2,8 @@
 from mmengine import ConfigDict
 
 from mmseg.models import build_segmentor
-from .utils import _segmentor_predict
+from tests.test_models.test_segmentors.utils import \
+    _segmentor_forward_train_test
 
 
 def test_multimodal_encoder_decoder():
@@ -20,5 +21,4 @@ def test_multimodal_encoder_decoder():
         train_cfg=None,
         test_cfg=dict(mode='whole'))
     segmentor = build_segmentor(cfg)
-    # test predict
-    _segmentor_predict(segmentor)
+    _segmentor_forward_train_test(segmentor)