| |
| from typing import Optional, Tuple, Union |
|
|
| import numpy as np |
| import torch |
| from mmengine.model import BaseModule |
| from torch import Tensor |
|
|
| from mmdet.registry import MODELS |
| from .mse_loss import mse_loss |
|
|
|
|
| @MODELS.register_module() |
| class MarginL2Loss(BaseModule): |
| """L2 loss with margin. |
| |
| Args: |
| neg_pos_ub (int, optional): The upper bound of negative to positive |
| samples in hard mining. Defaults to -1. |
| pos_margin (float, optional): The similarity margin for positive |
| samples in hard mining. Defaults to -1. |
| neg_margin (float, optional): The similarity margin for negative |
| samples in hard mining. Defaults to -1. |
| hard_mining (bool, optional): Whether to use hard mining. Defaults to |
| False. |
| reduction (str, optional): The method to reduce the loss. |
| Options are "none", "mean" and "sum". Defaults to "mean". |
| loss_weight (float, optional): The weight of loss. Defaults to 1.0. |
| """ |
|
|
| def __init__(self, |
| neg_pos_ub: int = -1, |
| pos_margin: float = -1, |
| neg_margin: float = -1, |
| hard_mining: bool = False, |
| reduction: str = 'mean', |
| loss_weight: float = 1.0): |
| super(MarginL2Loss, self).__init__() |
| self.neg_pos_ub = neg_pos_ub |
| self.pos_margin = pos_margin |
| self.neg_margin = neg_margin |
| self.hard_mining = hard_mining |
| self.reduction = reduction |
| self.loss_weight = loss_weight |
|
|
| def forward(self, |
| pred: Tensor, |
| target: Tensor, |
| weight: Optional[Tensor] = None, |
| avg_factor: Optional[float] = None, |
| reduction_override: Optional[str] = None) -> Tensor: |
| """Forward function. |
| |
| Args: |
| pred (torch.Tensor): The prediction. |
| target (torch.Tensor): The learning target of the prediction. |
| weight (torch.Tensor, optional): The weight of loss for each |
| prediction. Defaults to None. |
| avg_factor (float, optional): Average factor that is used to |
| average the loss. Defaults to None. |
| 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) |
| pred, weight, avg_factor = self.update_weight(pred, target, weight, |
| avg_factor) |
| loss_bbox = self.loss_weight * mse_loss( |
| pred, |
| target.float(), |
| weight.float(), |
| reduction=reduction, |
| avg_factor=avg_factor) |
| return loss_bbox |
|
|
| def update_weight(self, pred: Tensor, target: Tensor, weight: Tensor, |
| avg_factor: float) -> Tuple[Tensor, Tensor, float]: |
| """Update the weight according to targets. |
| |
| Args: |
| pred (torch.Tensor): The prediction. |
| target (torch.Tensor): The learning target of the prediction. |
| weight (torch.Tensor): The weight of loss for each prediction. |
| avg_factor (float): Average factor that is used to average the |
| loss. |
| |
| Returns: |
| tuple[torch.Tensor]: The updated prediction, weight and average |
| factor. |
| """ |
| if weight is None: |
| weight = target.new_ones(target.size()) |
|
|
| invalid_inds = weight <= 0 |
| target[invalid_inds] = -1 |
| pos_inds = target == 1 |
| neg_inds = target == 0 |
|
|
| if self.pos_margin > 0: |
| pred[pos_inds] -= self.pos_margin |
| if self.neg_margin > 0: |
| pred[neg_inds] -= self.neg_margin |
| pred = torch.clamp(pred, min=0, max=1) |
|
|
| num_pos = int((target == 1).sum()) |
| num_neg = int((target == 0).sum()) |
| if self.neg_pos_ub > 0 and num_neg / (num_pos + |
| 1e-6) > self.neg_pos_ub: |
| num_neg = num_pos * self.neg_pos_ub |
| neg_idx = torch.nonzero(target == 0, as_tuple=False) |
|
|
| if self.hard_mining: |
| costs = mse_loss( |
| pred, target.float(), |
| reduction='none')[neg_idx[:, 0], neg_idx[:, 1]].detach() |
| neg_idx = neg_idx[costs.topk(num_neg)[1], :] |
| else: |
| neg_idx = self.random_choice(neg_idx, num_neg) |
|
|
| new_neg_inds = neg_inds.new_zeros(neg_inds.size()).bool() |
| new_neg_inds[neg_idx[:, 0], neg_idx[:, 1]] = True |
|
|
| invalid_neg_inds = torch.logical_xor(neg_inds, new_neg_inds) |
| weight[invalid_neg_inds] = 0 |
|
|
| avg_factor = (weight > 0).sum() |
| return pred, weight, avg_factor |
|
|
| @staticmethod |
| def random_choice(gallery: Union[list, np.ndarray, Tensor], |
| num: int) -> np.ndarray: |
| """Random select some elements from the gallery. |
| |
| It seems that Pytorch's implementation is slower than numpy so we use |
| numpy to randperm the indices. |
| |
| Args: |
| gallery (list | np.ndarray | torch.Tensor): The gallery from |
| which to sample. |
| num (int): The number of elements to sample. |
| """ |
| assert len(gallery) >= num |
| if isinstance(gallery, list): |
| gallery = np.array(gallery) |
| cands = np.arange(len(gallery)) |
| np.random.shuffle(cands) |
| rand_inds = cands[:num] |
| if not isinstance(gallery, np.ndarray): |
| rand_inds = torch.from_numpy(rand_inds).long().to(gallery.device) |
| return gallery[rand_inds] |
|
|
