Add class weight

.github/workflows/build.yml

@@ -58,7 +58,7 @@ jobs:
       - name: Install PyTorch
         run: pip install torch==${{matrix.torch}}+cpu torchvision==${{matrix.torchvision}}+cpu -f https://download.pytorch.org/whl/torch_stable.html
       - name: Install MMCV
-        run: pip install mmcv-full==1.2.2 -f https://download.openmmlab.com/mmcv/dist/cpu/torch${{matrix.torch}}/index.html
+        run: pip install mmcv-full==1.2.4 -f https://download.openmmlab.com/mmcv/dist/cpu/torch${{matrix.torch}}/index.html
       - name: Install MMDet
         run: pip install -q git+https://github.com/open-mmlab/mmdetection/
       - name: Install unittest dependencies
@@ -138,7 +138,7 @@ jobs:
         run: pip install torch==${{matrix.torch}} torchvision==${{matrix.torchvision}} -f https://download.pytorch.org/whl/torch_stable.html
       - name: Install mmaction dependencies
         run: |
-          pip install mmcv-full==1.2.2 -f https://download.openmmlab.com/mmcv/dist/${{matrix.mmcv}}/index.html
+          pip install mmcv-full==1.2.4 -f https://download.openmmlab.com/mmcv/dist/${{matrix.mmcv}}/index.html
           pip install -q git+https://github.com/open-mmlab/mmdetection/
           pip install -r requirements.txt
       - name: Build and install

mmaction/models/losses/cross_entropy_loss.py

@@ -1,3 +1,4 @@
+import torch
 import torch.nn.functional as F
 
 from ..registry import LOSSES
@@ -6,7 +7,22 @@
 
 @LOSSES.register_module()
 class CrossEntropyLoss(BaseWeightedLoss):
-    """Cross Entropy Loss."""
+    """Cross Entropy Loss.
+
+    Args:
+        loss_weight (float): Factor scalar multiplied on the loss.
+            Default: 1.0.
+        class_weight (list[float] | None): Loss weight for each class. If set
+            as None, use the same weight 1 for all classes. Only applies
+            to CrossEntropyLoss and BCELossWithLogits (should not be set when
+            using other losses). Default: None.
+    """
+
+    def __init__(self, loss_weight=1.0, class_weight=None):
+        super().__init__(loss_weight=loss_weight)
+        self.class_weight = None
+        if class_weight is not None:
+            self.class_weight = torch.Tensor(class_weight)
 
     def _forward(self, cls_score, label, **kwargs):
         """Forward function.
@@ -20,13 +36,31 @@ def _forward(self, cls_score, label, **kwargs):
         Returns:
             torch.Tensor: The returned CrossEntropy loss.
         """
+        if self.class_weight is not None:
+            assert 'weight' not in kwargs, "The key 'weight' already exists."
+            kwargs['weight'] = self.class_weight.to(cls_score.device)
         loss_cls = F.cross_entropy(cls_score, label, **kwargs)
         return loss_cls
 
 
 @LOSSES.register_module()
 class BCELossWithLogits(BaseWeightedLoss):
-    """Binary Cross Entropy Loss with logits."""
+    """Binary Cross Entropy Loss with logits.
+
+    Args:
+        loss_weight (float): Factor scalar multiplied on the loss.
+            Default: 1.0.
+        class_weight (list[float] | None): Loss weight for each class. If set
+            as None, use the same weight 1 for all classes. Only applies
+            to CrossEntropyLoss and BCELossWithLogits (should not be set when
+            using other losses). Default: None.
+    """
+
+    def __init__(self, loss_weight=1.0, class_weight=None):
+        super().__init__(loss_weight=loss_weight)
+        self.class_weight = None
+        if class_weight is not None:
+            self.class_weight = torch.Tensor(class_weight)
 
     def _forward(self, cls_score, label, **kwargs):
         """Forward function.
@@ -40,6 +74,9 @@ def _forward(self, cls_score, label, **kwargs):
         Returns:
             torch.Tensor: The returned bce loss with logits.
         """
+        if self.class_weight is not None:
+            assert 'weight' not in kwargs, "The key 'weight' already exists."
+            kwargs['weight'] = self.class_weight.to(cls_score.device)
         loss_cls = F.binary_cross_entropy_with_logits(cls_score, label,
                                                       **kwargs)
         return loss_cls

tests/test_loss.py

@@ -81,21 +81,38 @@ def test_hvu_loss():
 
 def test_cross_entropy_loss():
     cls_scores = torch.rand((3, 4))
-    gt_labels = torch.LongTensor([2] * 3).squeeze()
+    gt_labels = torch.LongTensor([0, 1, 2]).squeeze()
 
     cross_entropy_loss = CrossEntropyLoss()
     output_loss = cross_entropy_loss(cls_scores, gt_labels)
     assert torch.equal(output_loss, F.cross_entropy(cls_scores, gt_labels))
 
+    weight = torch.rand(4)
+    class_weight = weight.numpy().tolist()
+    cross_entropy_loss = CrossEntropyLoss(class_weight=class_weight)
+    output_loss = cross_entropy_loss(cls_scores, gt_labels)
+    assert torch.equal(output_loss,
+                       F.cross_entropy(cls_scores, gt_labels, weight=weight))
+
 
 def test_bce_loss_with_logits():
     cls_scores = torch.rand((3, 4))
     gt_labels = torch.rand((3, 4))
+
     bce_loss_with_logits = BCELossWithLogits()
     output_loss = bce_loss_with_logits(cls_scores, gt_labels)
     assert torch.equal(
         output_loss, F.binary_cross_entropy_with_logits(cls_scores, gt_labels))
 
+    weight = torch.rand(4)
+    class_weight = weight.numpy().tolist()
+    bce_loss_with_logits = BCELossWithLogits(class_weight=class_weight)
+    output_loss = bce_loss_with_logits(cls_scores, gt_labels)
+    assert torch.equal(
+        output_loss,
+        F.binary_cross_entropy_with_logits(
+            cls_scores, gt_labels, weight=weight))
+
 
 def test_nll_loss():
     cls_scores = torch.randn(3, 3)