[Feature] Add MonoFlex Loss (#1114)

* add monoflex loss

* add unnitest

* add assert error

* fix typos

* remove monoflex_iou_loss

* remove monoflex_iou_loss unittest

mmdet3d/models/losses/__init__.py

@@ -2,11 +2,13 @@
 from mmdet.models.losses import FocalLoss, SmoothL1Loss, binary_cross_entropy
 from .axis_aligned_iou_loss import AxisAlignedIoULoss, axis_aligned_iou_loss
 from .chamfer_distance import ChamferDistance, chamfer_distance
+from .multibin_loss import MultiBinLoss
 from .paconv_regularization_loss import PAConvRegularizationLoss
 from .uncertain_smooth_l1_loss import UncertainL1Loss, UncertainSmoothL1Loss
 
 __all__ = [
     'FocalLoss', 'SmoothL1Loss', 'binary_cross_entropy', 'ChamferDistance',
     'chamfer_distance', 'axis_aligned_iou_loss', 'AxisAlignedIoULoss',
-    'PAConvRegularizationLoss', 'UncertainL1Loss', 'UncertainSmoothL1Loss'
+    'PAConvRegularizationLoss', 'UncertainL1Loss', 'UncertainSmoothL1Loss',
+    'MultiBinLoss'
 ]

mmdet3d/models/losses/multibin_loss.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from mmdet.models.builder import LOSSES
+from mmdet.models.losses.utils import weighted_loss
+
+
+@weighted_loss
+def multibin_loss(pred_orientations, gt_orientations, num_dir_bins=4):
+    """Multi-Bin Loss.
+
+    Args:
+        pred_orientations(torch.Tensor): Predicted local vector
+            orientation in [axis_cls, head_cls, sin, cos] format.
+            shape (N, num_dir_bins * 4)
+        gt_orientations(torch.Tensor): Corresponding gt bboxes,
+            shape (N, num_dir_bins * 2).
+        num_dir_bins(int, optional): Number of bins to encode
+            direction angle.
+            Defaults: 4.
+
+    Return:
+        torch.Tensor: Loss tensor.
+    """
+    cls_losses = 0
+    reg_losses = 0
+    reg_cnt = 0
+    for i in range(num_dir_bins):
+        # bin cls loss
+        cls_ce_loss = F.cross_entropy(
+            pred_orientations[:, (i * 2):(i * 2 + 2)],
+            gt_orientations[:, i].long(),
+            reduction='mean')
+        # regression loss
+        valid_mask_i = (gt_orientations[:, i] == 1)
+        cls_losses += cls_ce_loss
+        if valid_mask_i.sum() > 0:
+            start = num_dir_bins * 2 + i * 2
+            end = start + 2
+            pred_offset = F.normalize(pred_orientations[valid_mask_i,
+                                                        start:end])
+            gt_offset_sin = torch.sin(gt_orientations[valid_mask_i,
+                                                      num_dir_bins + i])
+            gt_offset_cos = torch.cos(gt_orientations[valid_mask_i,
+                                                      num_dir_bins + i])
+            reg_loss = \
+                F.l1_loss(pred_offset[:, 0], gt_offset_sin,
+                          reduction='none') + \
+                F.l1_loss(pred_offset[:, 1], gt_offset_cos,
+                          reduction='none')
+
+            reg_losses += reg_loss.sum()
+            reg_cnt += valid_mask_i.sum()
+
+        return cls_losses / num_dir_bins + reg_losses / reg_cnt
+
+
+@LOSSES.register_module()
+class MultiBinLoss(nn.Module):
+    """Multi-Bin Loss for orientation.
+
+    Args:
+        reduction (str, optional): The method to reduce the loss.
+            Options are 'none', 'mean' and 'sum'. Defaults to 'none'.
+        loss_weight (float, optional): The weight of loss. Defaults
+            to 1.0.
+    """
+
+    def __init__(self, reduction='none', loss_weight=1.0):
+        super(MultiBinLoss, self).__init__()
+        assert reduction in ['none', 'sum', 'mean']
+        self.reduction = reduction
+        self.loss_weight = loss_weight
+
+    def forward(self, pred, target, num_dir_bins, reduction_override=None):
+        """Forward function.
+
+        Args:
+            pred (torch.Tensor): The prediction.
+            target (torch.Tensor): The learning target of the prediction.
+            num_dir_bins (int): Number of bins to encode direction angle.
+            reduction_override (str, optional): The reduction method used to
+                override the original reduction method of the loss.
+                Defaults to None.
+        """
+        assert reduction_override in (None, 'none', 'mean', 'sum')
+        reduction = (
+            reduction_override if reduction_override else self.reduction)
+        loss = self.loss_weight * multibin_loss(
+            pred, target, num_dir_bins=num_dir_bins, reduction=reduction)
+        return loss

tests/test_metrics/test_losses.py

@@ -146,3 +146,29 @@ def test_uncertain_smooth_l1_loss():
     expected_smooth_l1_loss = torch.tensor(3.9795)
     assert torch.allclose(
         mean_smooth_l1_loss, expected_smooth_l1_loss, atol=1e-4)
+
+
+def test_multibin_loss():
+    from mmdet3d.models.losses import MultiBinLoss
+
+    # reduction should be in ['none', 'mean', 'sum']
+    with pytest.raises(AssertionError):
+        multibin_loss = MultiBinLoss(reduction='l2')
+
+    pred = torch.tensor([[
+        0.81, 0.32, 0.78, 0.52, 0.24, 0.12, 0.32, 0.11, 1.20, 1.30, 0.20, 0.11,
+        0.12, 0.11, 0.23, 0.31
+    ],
+                         [
+                             0.02, 0.19, 0.78, 0.22, 0.31, 0.12, 0.22, 0.11,
+                             1.20, 1.30, 0.45, 0.51, 0.12, 0.11, 0.13, 0.61
+                         ]])
+    target = torch.tensor([[1, 1, 0, 0, 2.14, 3.12, 0.68, -2.15],
+                           [1, 1, 0, 0, 3.12, 3.12, 2.34, 1.23]])
+    multibin_loss_cfg = dict(
+        type='MultiBinLoss', reduction='none', loss_weight=1.0)
+    multibin_loss = build_loss(multibin_loss_cfg)
+    output_multibin_loss = multibin_loss(pred, target, num_dir_bins=4)
+    expected_multibin_loss = torch.tensor(2.1120)
+    assert torch.allclose(
+        output_multibin_loss, expected_multibin_loss, atol=1e-4)