Reference for ultralytics/models/utils/loss.py

def forward(
    self,
    pred_bboxes: torch.Tensor,
    pred_scores: torch.Tensor,
    batch: Dict[str, Any],
    postfix: str = "",
    **kwargs: Any,
) -> Dict[str, torch.Tensor]:
    """
    Calculate loss for predicted bounding boxes and scores.

    Args:
        pred_bboxes (torch.Tensor): Predicted bounding boxes, shape (L, B, N, 4).
        pred_scores (torch.Tensor): Predicted class scores, shape (L, B, N, C).
        batch (Dict[str, Any]): Batch information containing cls, bboxes, and gt_groups.
        postfix (str, optional): Postfix for loss names.
        **kwargs (Any): Additional arguments, may include 'match_indices'.

    Returns:
        (Dict[str, torch.Tensor]): Computed losses, including main and auxiliary (if enabled).

    Notes:
        Uses last elements of pred_bboxes and pred_scores for main loss, and the rest for auxiliary losses if
        self.aux_loss is True.
    """
    self.device = pred_bboxes.device
    match_indices = kwargs.get("match_indices", None)
    gt_cls, gt_bboxes, gt_groups = batch["cls"], batch["bboxes"], batch["gt_groups"]

    total_loss = self._get_loss(
        pred_bboxes[-1], pred_scores[-1], gt_bboxes, gt_cls, gt_groups, postfix=postfix, match_indices=match_indices
    )

    if self.aux_loss:
        total_loss.update(
            self._get_loss_aux(
                pred_bboxes[:-1], pred_scores[:-1], gt_bboxes, gt_cls, gt_groups, match_indices, postfix
            )
        )

    return total_loss

def forward(
    self,
    preds: Tuple[torch.Tensor, torch.Tensor],
    batch: Dict[str, Any],
    dn_bboxes: Optional[torch.Tensor] = None,
    dn_scores: Optional[torch.Tensor] = None,
    dn_meta: Optional[Dict[str, Any]] = None,
) -> Dict[str, torch.Tensor]:
    """
    Forward pass to compute detection loss with optional denoising loss.

    Args:
        preds (Tuple[torch.Tensor, torch.Tensor]): Tuple containing predicted bounding boxes and scores.
        batch (Dict[str, Any]): Batch data containing ground truth information.
        dn_bboxes (torch.Tensor, optional): Denoising bounding boxes.
        dn_scores (torch.Tensor, optional): Denoising scores.
        dn_meta (Dict[str, Any], optional): Metadata for denoising.

    Returns:
        (Dict[str, torch.Tensor]): Dictionary containing total loss and denoising loss if applicable.
    """
    pred_bboxes, pred_scores = preds
    total_loss = super().forward(pred_bboxes, pred_scores, batch)

    # Check for denoising metadata to compute denoising training loss
    if dn_meta is not None:
        dn_pos_idx, dn_num_group = dn_meta["dn_pos_idx"], dn_meta["dn_num_group"]
        assert len(batch["gt_groups"]) == len(dn_pos_idx)

        # Get the match indices for denoising
        match_indices = self.get_dn_match_indices(dn_pos_idx, dn_num_group, batch["gt_groups"])

        # Compute the denoising training loss
        dn_loss = super().forward(dn_bboxes, dn_scores, batch, postfix="_dn", match_indices=match_indices)
        total_loss.update(dn_loss)
    else:
        # If no denoising metadata is provided, set denoising loss to zero
        total_loss.update({f"{k}_dn": torch.tensor(0.0, device=self.device) for k in total_loss.keys()})

    return total_loss

@staticmethod
def get_dn_match_indices(
    dn_pos_idx: List[torch.Tensor], dn_num_group: int, gt_groups: List[int]
) -> List[Tuple[torch.Tensor, torch.Tensor]]:
    """
    Get match indices for denoising.

    Args:
        dn_pos_idx (List[torch.Tensor]): List of tensors containing positive indices for denoising.
        dn_num_group (int): Number of denoising groups.
        gt_groups (List[int]): List of integers representing number of ground truths per image.

    Returns:
        (List[Tuple[torch.Tensor, torch.Tensor]]): List of tuples containing matched indices for denoising.
    """
    dn_match_indices = []
    idx_groups = torch.as_tensor([0, *gt_groups[:-1]]).cumsum_(0)
    for i, num_gt in enumerate(gt_groups):
        if num_gt > 0:
            gt_idx = torch.arange(end=num_gt, dtype=torch.long) + idx_groups[i]
            gt_idx = gt_idx.repeat(dn_num_group)
            assert len(dn_pos_idx[i]) == len(gt_idx), (
                f"Expected the same length, but got {len(dn_pos_idx[i])} and {len(gt_idx)} respectively."
            )
            dn_match_indices.append((dn_pos_idx[i], gt_idx))
        else:
            dn_match_indices.append((torch.zeros([0], dtype=torch.long), torch.zeros([0], dtype=torch.long)))
    return dn_match_indices