[Feature] Add mask channel in MGD Loss

mmrazor/models/architectures/connectors/mgd_connector.py

@@ -27,10 +27,12 @@ def __init__(
         student_channels: int,
         teacher_channels: int,
         lambda_mgd: float = 0.65,
+        mask_on_channel: bool = False,
         init_cfg: Optional[Dict] = None,
     ) -> None:
         super().__init__(init_cfg)
         self.lambda_mgd = lambda_mgd
+        self.mask_on_channel = mask_on_channel
         if student_channels != teacher_channels:
             self.align = nn.Conv2d(
                 student_channels,
@@ -55,7 +57,11 @@ def forward_train(self, feature: torch.Tensor) -> torch.Tensor:
         N, C, H, W = feature.shape
 
         device = feature.device
-        mat = torch.rand((N, 1, H, W)).to(device)
+        if not self.mask_on_channel:
+            mat = torch.rand((N, 1, H, W)).to(device)
+        else:
+            mat = torch.rand((N, C, 1, 1)).to(device)
+
         mat = torch.where(mat > 1 - self.lambda_mgd,
                           torch.zeros(1).to(device),
                           torch.ones(1).to(device)).to(device)

tests/test_models/test_architectures/test_connectors/test_connectors.py

@@ -143,3 +143,19 @@ def test_mgd_connector(self):
 
         assert s_output1.shape == torch.Size([1, 16, 8, 8])
         assert s_output2.shape == torch.Size([1, 32, 8, 8])
+
+        mgd_connector1 = MGDConnector(
+            student_channels=16,
+            teacher_channels=16,
+            lambda_mgd=0.65,
+            mask_on_channel=True)
+        mgd_connector2 = MGDConnector(
+            student_channels=16,
+            teacher_channels=32,
+            lambda_mgd=0.65,
+            mask_on_channel=True)
+        s_output1 = mgd_connector1.forward_train(s_feat)
+        s_output2 = mgd_connector2.forward_train(s_feat)
+
+        assert s_output1.shape == torch.Size([1, 16, 8, 8])
+        assert s_output2.shape == torch.Size([1, 32, 8, 8])