API Reference: Losses & Assigners

Task-specific loss functions and label assigners.

---

Loss Functions

Classification Losses

corecv.losses.classification

GPU-native classification loss functions for CoreCV.

Provides focal loss and label-smoothing cross-entropy, both implemented as pure vectorised PyTorch operations with zero CPU-GPU synchronisations during forward or backward passes. No Python for-loops, no .item() calls, no .cpu() transfers — everything stays on VRAM.

Example

import torch from corecv.losses.classification import FocalLoss loss_fn = FocalLoss(alpha=0.25, gamma=2.0) logits = torch.randn(8, 80, 32, 32, device="cuda") targets = torch.randint(0, 80, (8, 32, 32), device="cuda") loss = loss_fn(logits, targets)

FocalLoss(alpha=0.25, gamma=2.0, reduction='mean')

Bases: Module

Focal Loss for dense object detection (Lin et al. ICCV 2017).

Reduces the relative loss for well-classified examples, focusing training on hard negatives. The formulation is::

FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)

where p_t = exp(-CE) and alpha_t is a per-class weighting factor.

This implementation computes the standard cross-entropy with reduction='none' and applies the focal modulator entirely via vectorised tensor ops — no Python loops.

Parameters:

Name	Type	Description	Default
alpha	float | Tensor	Weighting factor. Can be a scalar float applied to all classes, or a 1-D Tensor of shape (C,) with per-class weights.	0.25
gamma	float	Focal focusing parameter. Higher values down-weight easy examples more aggressively. Default 2.0.	2.0
reduction	str	Aggregation mode. 'mean' (default) returns the mean loss over all elements; 'sum' returns the sum; 'none' returns the raw per-element loss tensor.	'mean'

Example

loss_fn = FocalLoss(alpha=0.25, gamma=2.0, reduction="mean") logits = torch.randn(4, 80, 16, 16, device="cuda") targets = torch.randint(0, 80, (4, 16, 16), device="cuda") loss = loss_fn(logits, targets)

Initialise FocalLoss with alpha, gamma, and reduction mode.

Source code in src/corecv/losses/classification.py

def __init__(
    self,
    alpha: float | Tensor = 0.25,
    gamma: float = 2.0,
    reduction: str = "mean",
) -> None:
    """Initialise FocalLoss with alpha, gamma, and reduction mode."""
    super().__init__()
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}. Must be 'mean', 'sum', or 'none'."
        raise ValueError(msg)

    self.gamma = float(gamma)
    self.reduction = reduction

    # Store alpha as a buffer so it moves with .to(device) automatically.
    # If a Tensor is passed it is registered directly; a scalar float is
    # broadcast later during forward.
    if isinstance(alpha, Tensor):
        self.register_buffer("alpha", alpha.float())
    else:
        # Will be broadcast to match logits shape in forward
        self.alpha = float(alpha)

forward(inputs, targets)

Compute focal loss.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Unnormalised logits of shape (B, C, ...) where C is the number of classes. Any trailing spatial dimensions are supported.	required
targets	Tensor	Ground-truth class indices of shape (B, ...) with the same trailing spatial dimensions as inputs (excluding the class axis). Values must be in [0, C).	required

Returns:

Type	Description
Tensor	Scalar or tensor loss depending on self.reduction.

Source code in src/corecv/losses/classification.py

def forward(self, inputs: Tensor, targets: Tensor) -> Tensor:
    """Compute focal loss.

    Args:
        inputs: Unnormalised logits of shape ``(B, C, ...)`` where ``C``
            is the number of classes.  Any trailing spatial dimensions
            are supported.
        targets: Ground-truth class indices of shape ``(B, ...)`` with
            the same trailing spatial dimensions as *inputs* (excluding
            the class axis).  Values must be in ``[0, C)``.

    Returns:
        Scalar or tensor loss depending on ``self.reduction``.
    """
    # Standard cross-entropy without reduction — gives per-element loss
    ce_loss: Tensor = F.cross_entropy(
        inputs,
        targets,
        reduction="none",
    )

    # p_t = exp(-CE) is the model's estimated probability for the
    # ground-truth class.
    pt: Tensor = torch.exp(-ce_loss)

    # Focal modulator: (1 - p_t)^gamma
    focal_weight: Tensor = (1.0 - pt).pow(self.gamma)

    # Alpha weighting: broadcast scalar or match channel dim
    alpha = self.alpha
    if isinstance(alpha, float):
        # Scalar — simply multiply
        loss = alpha * focal_weight * ce_loss
    else:
        # Per-class alpha tensor of shape (C,).
        # Gather the alpha value for each spatial position's target
        # class.  We expand alpha to full (B, C, *spatial) shape,
        # then use torch.gather along dim=1 with the target indices.
        alpha_shape = [1] * inputs.dim()  # (1, C, 1, 1, ...)
        alpha_shape[1] = -1  # sentinel for C
        alpha_expanded: Tensor = alpha.view(alpha_shape)
        # Expand to full (B, C, *spatial) — expand does not allocate
        alpha_expanded = alpha_expanded.expand_as(inputs)
        # Gather along class dim using target indices
        alpha_t: Tensor = torch.gather(
            alpha_expanded,
            dim=1,
            index=targets.unsqueeze(1),
        ).squeeze(1)
        loss = alpha_t * focal_weight * ce_loss

    return self._reduce(loss)

LabelSmoothingCrossEntropy(smoothing=0.1, reduction='mean')

Bases: Module

Label-smoothing cross-entropy loss.

Instead of a hard one-hot target [0, 0, 1, 0, ...], this loss assigns (1 - smoothing) probability mass to the correct class and distributes smoothing / (K - 1) mass uniformly across all other classes, where K is the number of classes. This discourages the model from becoming over-confident and acts as a regulariser.

Formulation::

L = (1 - smoothing) * CE(y, p) + smoothing * mean_k(log_softmax(z))

where CE is the standard cross-entropy and mean_k averages the log-softmax over all K classes.

The implementation is fully vectorised using F.log_softmax + F.nll_loss with no Python loops and no CPU-GPU sync points.

Parameters:

Name	Type	Description	Default
smoothing	float	Label-smoothing factor in [0, 1). 0.0 is equivalent to standard cross-entropy. Default 0.1.	0.1
reduction	str	Aggregation mode ('mean' | 'sum' | 'none'). Default 'mean'.	'mean'

Example

loss_fn = LabelSmoothingCrossEntropy(smoothing=0.1) logits = torch.randn(16, 100, device="cuda") targets = torch.randint(0, 100, (16,), device="cuda") loss = loss_fn(logits, targets)

Initialise LabelSmoothingCrossEntropy with smoothing factor and reduction mode.

Source code in src/corecv/losses/classification.py

def __init__(self, smoothing: float = 0.1, reduction: str = "mean") -> None:
    """Initialise LabelSmoothingCrossEntropy with smoothing factor and reduction mode."""
    super().__init__()
    if not (0.0 <= smoothing < 1.0):
        msg = f"smoothing must be in [0, 1), got {smoothing}"
        raise ValueError(msg)
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}. Must be 'mean', 'sum', or 'none'."
        raise ValueError(msg)

    self.smoothing = float(smoothing)
    self.reduction = reduction

forward(inputs, targets)

Compute label-smoothing cross-entropy.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Unnormalised logits of shape (B, C, ...) where C is the number of classes.	required
targets	Tensor	Ground-truth class indices of shape (B, ...) with values in [0, C).	required

Returns:

Type	Description
Tensor	Scalar or tensor loss depending on self.reduction.

Source code in src/corecv/losses/classification.py

def forward(self, inputs: Tensor, targets: Tensor) -> Tensor:
    """Compute label-smoothing cross-entropy.

    Args:
        inputs: Unnormalised logits of shape ``(B, C, ...)`` where ``C``
            is the number of classes.
        targets: Ground-truth class indices of shape ``(B, ...)`` with
            values in ``[0, C)``.

    Returns:
        Scalar or tensor loss depending on ``self.reduction``.
    """
    log_probs: Tensor = F.log_softmax(inputs, dim=1)

    # NLL part: -log_probs gathered at target positions, averaged over
    # all spatial / batch dimensions.
    nll_loss: Tensor = F.nll_loss(
        log_probs,
        targets,
        reduction="none",
    )

    # Smooth part: negative mean of log-softmax over all classes.
    # log_probs has shape (B, C, ...); mean over dim=1 gives (B, ...).
    smooth_loss: Tensor = -log_probs.mean(dim=1)

    # Combine
    loss: Tensor = (
        (1.0 - self.smoothing) * nll_loss + self.smoothing * smooth_loss
    )

    return self._reduce(loss)

Detection Losses

corecv.losses.detection

GPU-native detection loss functions for CoreCV.

Provides GIoU, CIoU, Quality Focal Loss and Varifocal Loss — all implemented as pure vectorised PyTorch operations with zero CPU-GPU synchronisations during forward or backward passes. No Python loops, no .item() calls, no .cpu() transfers — every computation remains on VRAM.

Supports both per-level tensors (B, 4, H, W) and flattened tensors (N, 4) for detection boxes. All IoU losses use torch.min / torch.max / torch.clamp for vectorised intersection / union / enclosing-box computation.

Example

import torch from corecv.losses.detection import GIoULoss, CIoULoss giou = GIoULoss(reduction="mean") ciou = CIoULoss(reduction="mean") pred = torch.randn(4, 4, 32, 32, device="cuda").abs() * 100 tgt = torch.randn(4, 4, 32, 32, device="cuda").abs() * 100 loss = giou(pred, tgt)

CIoULoss(reduction='mean', eps=1e-07)

Bases: Module

Complete IoU Loss (Zheng et al. CVPR 2020).

Extends DIoU with an aspect-ratio consistency term::

CIoU = IoU - rho²(b_pred, b_gt) / c² - alpha * v

where:

• rho² is the squared centre distance.
• c is the diagonal of the enclosing box.
• v measures aspect-ratio similarity: v = (4 / pi²) * (arctan(w_gt/h_gt) - arctan(w_pred/h_pred))²
• alpha = v / (1 - IoU + v) balances the two penalty terms.

The loss is 1 - CIoU.

Fully vectorised: all operations run on batched tensors with no Python loops and no CPU-GPU synchronisations.

Parameters:

Name	Type	Description	Default
reduction	str	'mean' | 'sum' | 'none'. Default 'mean'.	'mean'
eps	float	Small constant for numerical stability. Default 1e-7.	1e-07

Example

ciou = CIoULoss(reduction="mean") pred = torch.randn(4, 4, 32, 32, device="cuda").abs() * 100 tgt = torch.randn(4, 4, 32, 32, device="cuda").abs() * 100 loss = ciou(pred, tgt)

Initialise CIoULoss with reduction mode and epsilon.

Source code in src/corecv/losses/detection.py

def __init__(self, reduction: str = "mean", eps: float = 1e-7) -> None:
    """Initialise CIoULoss with reduction mode and epsilon."""
    super().__init__()
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}"
        raise ValueError(msg)
    self.reduction = reduction
    self.eps = eps

forward(pred_boxes, target_boxes)

Compute CIoU loss.

Parameters:

Name	Type	Description	Default
pred_boxes	Tensor	Predicted boxes, (B, 4, H, W) or (N, 4) in (l, t, r, b) format.	required
target_boxes	Tensor	Target boxes, same shape.	required

Returns:

Type	Description
Tensor	Scalar or tensor loss.

Source code in src/corecv/losses/detection.py

def forward(
    self,
    pred_boxes: Tensor,
    target_boxes: Tensor,
) -> Tensor:
    """Compute CIoU loss.

    Args:
        pred_boxes: Predicted boxes, ``(B, 4, H, W)`` or ``(N, 4)``
            in ``(l, t, r, b)`` format.
        target_boxes: Target boxes, same shape.

    Returns:
        Scalar or tensor loss.
    """
    # Flatten to (M, 4)
    if pred_boxes.dim() == _DIM_4D:
        M = pred_boxes.shape[0] * pred_boxes.shape[2] * pred_boxes.shape[3]
        pred_flat = pred_boxes.permute(0, 2, 3, 1).reshape(M, 4)
        target_flat = target_boxes.permute(0, 2, 3, 1).reshape(M, 4)
    elif pred_boxes.dim() == _DIM_2D:
        pred_flat = pred_boxes
        target_flat = target_boxes
    else:
        msg = f"Expected 2-D or 4-D boxes, got {pred_boxes.dim()}-D"
        raise ValueError(msg)

    # ---- Compute intersection directly for IoU ----------------------
    inter_l = torch.max(pred_flat[..., 0], target_flat[..., 0])
    inter_t = torch.max(pred_flat[..., 1], target_flat[..., 1])
    inter_r = torch.min(pred_flat[..., 2], target_flat[..., 2])
    inter_b = torch.min(pred_flat[..., 3], target_flat[..., 3])

    inter_area: Tensor = (inter_r - inter_l).clamp(min=0) * (
        inter_b - inter_t
    ).clamp(min=0)

    area_pred: Tensor = _box_area(pred_flat)
    area_target: Tensor = _box_area(target_flat)
    union: Tensor = area_pred + area_target - inter_area

    iou: Tensor = inter_area / (union + self.eps)

    # ---- Centre distance penalty ------------------------------------
    centre_dist2: Tensor = _centre_distance_squared(pred_flat, target_flat)
    enclose: Tensor = _enclosing_box(pred_flat, target_flat)
    diag2: Tensor = _diagonal_squared(enclose)
    # Clamp diagonal to prevent division-by-near-zero on degenerate
    # enclosing boxes (e.g. when input boxes have l > r or t > b).
    diag2 = diag2.clamp(min=1.0)
    rho2: Tensor = centre_dist2 / diag2

    # ---- Aspect-ratio consistency term -------------------------------
    pred_w = (pred_flat[..., 2] - pred_flat[..., 0]).abs().clamp(min=self.eps)
    pred_h = (pred_flat[..., 3] - pred_flat[..., 1]).abs().clamp(min=self.eps)
    tgt_w = (target_flat[..., 2] - target_flat[..., 0]).abs().clamp(min=self.eps)
    tgt_h = (target_flat[..., 3] - target_flat[..., 1]).abs().clamp(min=self.eps)

    v: Tensor = (4.0 / (torch.pi ** 2)) * (
        torch.atan(tgt_w / tgt_h) - torch.atan(pred_w / pred_h)
    ).pow(2)

    alpha: Tensor = v / (1.0 - iou + v + self.eps)

    # ---- CIoU loss ---------------------------------------------------
    ciou: Tensor = iou - rho2 - alpha * v
    loss: Tensor = (1.0 - ciou).clamp(max=100.0)

    return self._reduce(loss)

GIoULoss(reduction='mean', eps=1e-07)

Bases: Module

Generalized IoU Loss (Rezatofighi et al. CVPR 2019).

Extends IoU with a penalty term for non-overlapping boxes::

GIoU = IoU - (C - U) / C

where C is the enclosing-box area and U is the union area. The loss is 1 - GIoU.

Fully vectorised: all operations use torch.min / torch.max / torch.clamp on batched tensors with no Python loops.

Supports two input formats:

• Per-level: (B, 4, H, W) — typical output of anchor-free detection heads at a single FPN level.
• Flattened: (N, 4) — all predictions from a single level (or all levels concatenated).

When per_sample is True (default) and the input is 4-D, the loss is computed element-wise and returned with the requested reduction.

Parameters:

Name	Type	Description	Default
reduction	str	'mean' | 'sum' | 'none'. Default 'mean'.	'mean'
eps	float	Small constant for numerical stability. Default 1e-7.	1e-07

Example

giou = GIoULoss(reduction="mean") pred = torch.randn(4, 4, 32, 32, device="cuda").abs() tgt = torch.randn(4, 4, 32, 32, device="cuda").abs() loss = giou(pred, tgt)

Initialise GIoULoss with reduction mode and epsilon.

Source code in src/corecv/losses/detection.py

def __init__(self, reduction: str = "mean", eps: float = 1e-7) -> None:
    """Initialise GIoULoss with reduction mode and epsilon."""
    super().__init__()
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}"
        raise ValueError(msg)
    self.reduction = reduction
    self.eps = eps

forward(pred_boxes, target_boxes)

Compute GIoU loss.

Parameters:

Name	Type	Description	Default
pred_boxes	Tensor	Predicted boxes. Shape (B, 4, H, W) or (N, 4) in (l, t, r, b) format.	required
target_boxes	Tensor	Target boxes. Same shape as pred_boxes.	required

Returns:

Type	Description
Tensor	Scalar or tensor loss depending on self.reduction.

Source code in src/corecv/losses/detection.py

def forward(
    self,
    pred_boxes: Tensor,
    target_boxes: Tensor,
) -> Tensor:
    """Compute GIoU loss.

    Args:
        pred_boxes: Predicted boxes.  Shape ``(B, 4, H, W)`` or
            ``(N, 4)`` in ``(l, t, r, b)`` format.
        target_boxes: Target boxes.  Same shape as *pred_boxes*.

    Returns:
        Scalar or tensor loss depending on ``self.reduction``.
    """
    # Flatten to (M, 4) for uniform processing
    if pred_boxes.dim() == _DIM_4D:
        # (B, 4, H, W) -> (B*H*W, 4) — treat each spatial cell as a box
        M = pred_boxes.shape[0] * pred_boxes.shape[2] * pred_boxes.shape[3]
        pred_flat = pred_boxes.permute(0, 2, 3, 1).reshape(M, 4)
        target_flat = target_boxes.permute(0, 2, 3, 1).reshape(M, 4)
    elif pred_boxes.dim() == _DIM_2D:
        M = pred_boxes.shape[0]
        pred_flat = pred_boxes
        target_flat = target_boxes
    else:
        msg = f"Expected 2-D or 4-D boxes, got {pred_boxes.dim()}-D"
        raise ValueError(msg)

    # ---- Compute intersection / union directly from coordinates ------
    # This avoids the circular dependency of recovering union from IoU.
    inter_l = torch.max(pred_flat[..., 0], target_flat[..., 0])
    inter_t = torch.max(pred_flat[..., 1], target_flat[..., 1])
    inter_r = torch.min(pred_flat[..., 2], target_flat[..., 2])
    inter_b = torch.min(pred_flat[..., 3], target_flat[..., 3])

    inter_area: Tensor = (inter_r - inter_l).clamp(min=0) * (
        inter_b - inter_t
    ).clamp(min=0)

    area_pred: Tensor = _box_area(pred_flat)
    area_target: Tensor = _box_area(target_flat)
    union: Tensor = area_pred + area_target - inter_area

    iou: Tensor = inter_area / (union + self.eps)

    # Enclosing box
    enclose: Tensor = _enclosing_box(pred_flat, target_flat)
    enclose_area: Tensor = _box_area(enclose)

    # GIoU
    giou: Tensor = iou - (enclose_area - union) / (enclose_area + self.eps)
    loss: Tensor = 1.0 - giou

    return self._reduce(loss)

QualityFocalLoss(beta=2.0, reduction='mean')

Bases: Module

Quality Focal Loss (QFL) for anchor-free detection heads.

QFL replaces the standard binary cross-entropy with a modulated loss that treats the target as a continuous IoU quality score rather than a hard binary label::

loss(q) = |pred - q|^beta * BCE(pred, q)

where q is the IoU between the predicted and ground-truth box for positive samples (and 0 for negatives).

This formulation encourages the classification branch to directly predict IoU quality, enabling better ranking at inference.

Fully vectorised: uses F.binary_cross_entropy_with_logits with a pos_weight tensor and the quality-based weighting computed via pure tensor ops.

Parameters:

Name	Type	Description	Default
beta	float	Focusing parameter of the quality modulation. Default 2.0.	2.0
reduction	str	'mean' | 'sum' | 'none'. Default 'mean'.	'mean'

Example

qfl = QualityFocalLoss(beta=2.0) pred = torch.randn(1000, 80, device="cuda") tgt_scores = torch.rand(1000, 80, device="cuda") # IoU quality tgt_labels = torch.randint(0, 80, (1000,), device="cuda") loss = qfl(pred, tgt_scores, tgt_labels)

Initialise QualityFocalLoss with beta and reduction mode.

Source code in src/corecv/losses/detection.py

def __init__(self, beta: float = 2.0, reduction: str = "mean") -> None:
    """Initialise QualityFocalLoss with beta and reduction mode."""
    super().__init__()
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}"
        raise ValueError(msg)
    self.beta = float(beta)
    self.reduction = reduction

forward(pred_scores, target_scores, target_labels)

Compute Quality Focal Loss.

Parameters:

Name	Type	Description	Default
pred_scores	Tensor	Predicted classification logits. Shape (N, C) or (B, C, H, W).	required
target_scores	Tensor	Continuous quality targets (e.g. IoU values). Same shape as pred_scores. Values in [0, 1].	required
target_labels	Tensor	Integer class labels of shape (N,) or (B, H, W). Values in [0, C). Only used to determine which channel each sample belongs to.	required

Returns:

Type	Description
Tensor	Scalar or tensor loss.

Source code in src/corecv/losses/detection.py

def forward(
    self,
    pred_scores: Tensor,
    target_scores: Tensor,
    target_labels: Tensor,
) -> Tensor:
    """Compute Quality Focal Loss.

    Args:
        pred_scores: Predicted classification logits.
            Shape ``(N, C)`` or ``(B, C, H, W)``.
        target_scores: Continuous quality targets (e.g. IoU values).
            Same shape as *pred_scores*.  Values in ``[0, 1]``.
        target_labels: Integer class labels of shape ``(N,)`` or
            ``(B, H, W)``.  Values in ``[0, C)``.  Only used to
            determine which channel each sample belongs to.

    Returns:
        Scalar or tensor loss.
    """
    # Flatten if 4-D spatial format
    if pred_scores.dim() == _DIM_4D:
        B, C, H, W = pred_scores.shape
        pred_scores = pred_scores.permute(0, 2, 3, 1).reshape(B * H * W, C)
        target_scores = target_scores.permute(0, 2, 3, 1).reshape(B * H * W, C)
        target_labels = target_labels.reshape(B * H * W)

    # One-hot mask for the target class channel: (N, C)
    num_classes = pred_scores.shape[1]
    one_hot: Tensor = F.one_hot(target_labels, num_classes=num_classes).float()

    # Binary cross-entropy per element (unreduced)
    bce: Tensor = F.binary_cross_entropy_with_logits(
        pred_scores, target_scores, reduction="none"
    )

    # Quality modulation: |pred - q|^beta
    # sigmoid(pred) - target  (both in [0, 1] range)
    pred_prob: Tensor = pred_scores.sigmoid()
    modulator: Tensor = (pred_prob - target_scores.detach()).abs().pow(self.beta)

    # Only the target class contributes to the loss
    loss: Tensor = modulator * bce * one_hot

    # Sum over classes, mean over samples
    loss_per_sample: Tensor = loss.sum(dim=1)

    return self._reduce(loss_per_sample)

VarifocalLoss(gamma=2.0, alpha=0.25, reduction='mean')

Bases: Module

Varifocal Loss (VFL) for dense object detection.

An asymmetric focal loss that re-weights positive and negative samples differently based on the quality (IoU) of each sample::

VFL(p, q) = -q * log(p)              if positive
            = -alpha * p^gamma * log(1-p)   if negative

where q is the IoU quality score (target) for positives, p is the predicted probability, gamma is the focusing parameter, and alpha balances positive / negative contribution.

This encourages the model to predict high scores only for high-quality detections while suppressing low-quality false positives.

Fully vectorised: all operations are on batched tensors with no Python loops.

Parameters:

Name	Type	Description	Default
gamma	float	Focusing parameter for negative samples. Default 2.0.	2.0
alpha	float	Balancing factor for negative samples. Default 0.25.	0.25
reduction	str	'mean' | 'sum' | 'none'. Default 'mean'.	'mean'

Example

vfl = VarifocalLoss(gamma=2.0, alpha=0.25) pred = torch.randn(500, 80, device="cuda") tgt_scores = torch.rand(500, 80, device="cuda") tgt_labels = torch.randint(0, 80, (500,), device="cuda") loss = vfl(pred, tgt_scores, tgt_labels)

Initialise VarifocalLoss with gamma, alpha, and reduction mode.

Source code in src/corecv/losses/detection.py

def __init__(
    self,
    gamma: float = 2.0,
    alpha: float = 0.25,
    reduction: str = "mean",
) -> None:
    """Initialise VarifocalLoss with gamma, alpha, and reduction mode."""
    super().__init__()
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}"
        raise ValueError(msg)
    self.gamma = float(gamma)
    self.alpha = float(alpha)
    self.reduction = reduction

forward(pred_scores, target_scores, target_labels)

Compute Varifocal Loss.

Parameters:

Name	Type	Description	Default
pred_scores	Tensor	Predicted classification logits, (N, C) or (B, C, H, W).	required
target_scores	Tensor	Quality targets in [0, 1], same shape as pred_scores. Positive positions carry their IoU value; negative positions carry 0.	required
target_labels	Tensor	Integer class labels, (N,) or (B, H, W).	required

Returns:

Type	Description
Tensor	Scalar or tensor loss.

Source code in src/corecv/losses/detection.py

def forward(
    self,
    pred_scores: Tensor,
    target_scores: Tensor,
    target_labels: Tensor,
) -> Tensor:
    """Compute Varifocal Loss.

    Args:
        pred_scores: Predicted classification logits, ``(N, C)`` or
            ``(B, C, H, W)``.
        target_scores: Quality targets in ``[0, 1]``, same shape as
            *pred_scores*.  Positive positions carry their IoU value;
            negative positions carry ``0``.
        target_labels: Integer class labels, ``(N,)`` or
            ``(B, H, W)``.

    Returns:
        Scalar or tensor loss.
    """
    # Flatten if 4-D
    if pred_scores.dim() == _DIM_4D:
        B, C, H, W = pred_scores.shape
        pred_scores = pred_scores.permute(0, 2, 3, 1).reshape(B * H * W, C)
        target_scores = target_scores.permute(0, 2, 3, 1).reshape(B * H * W, C)
        target_labels = target_labels.reshape(B * H * W)

    num_classes = pred_scores.shape[1]
    one_hot: Tensor = F.one_hot(target_labels, num_classes=num_classes).float()

    # Predicted probability
    pred_prob: Tensor = pred_scores.sigmoid()

    # ---- Positive loss: -q * log(p) ---------------------------------
    # Only for the target class; zeros elsewhere
    pos_loss: Tensor = -target_scores * torch.log(pred_prob.clamp(min=1e-8))

    # ---- Negative loss: -alpha * p^gamma * log(1-p) ------------------
    neg_loss: Tensor = -self.alpha * pred_prob.pow(self.gamma) * torch.log(
        (1.0 - pred_prob).clamp(min=1e-8)
    )

    # Combine using masks
    combined: Tensor = pos_loss + neg_loss
    loss_per_class: Tensor = combined * one_hot  # (N, C)

    # Sum over classes, keep per-sample: (N,)
    loss_per_sample: Tensor = loss_per_class.sum(dim=1)

    return self._reduce(loss_per_sample)

Segmentation Losses

corecv.losses.segmentation

GPU-native segmentation loss functions for CoreCV.

Provides Dice loss and a combined CrossEntropy + Dice loss, both implemented as pure vectorised PyTorch operations with zero CPU-GPU synchronisations during forward or backward passes. No Python for-loops over channels or batch elements — every computation uses broadcasting, F.one_hot, and standard tensor ops that remain entirely on VRAM.

Example

import torch from corecv.losses.segmentation import DiceLoss loss_fn = DiceLoss(smooth=1.0) logits = torch.randn(4, 21, 256, 256, device="cuda") targets = torch.randint(0, 21, (4, 256, 256), device="cuda") loss = loss_fn(logits, targets)

CombinedSegmentationLoss(ce_weight=1.0, dice_weight=1.0, ignore_index=-100)

Bases: Module

Combined Cross-Entropy + Dice loss for segmentation.

Summation of a weighted F.cross_entropy term and a :class:DiceLoss term, designed for common semantic segmentation training where both pixel-wise classification accuracy and region overlap matter.

Handles ignore_index transparently: pixels marked with ignore_index in the target are excluded from both the CE and the Dice computation, ensuring no gradient leakage from unlabeled regions.

The entire pipeline is 100% GPU-native with zero CPU-GPU sync points.

Parameters:

Name	Type	Description	Default
ce_weight	float	Weight of the cross-entropy term. Default 1.0.	1.0
dice_weight	float	Weight of the Dice term. Default 1.0.	1.0
ignore_index	int	Class index to ignore during loss computation. Default -100 (same as F.cross_entropy default).	-100

Example

loss_fn = CombinedSegmentationLoss( ... ce_weight=1.0, dice_weight=1.0, ignore_index=255, ... ) logits = torch.randn(8, 21, 256, 256, device="cuda") targets = torch.full((8, 256, 256), 255, device="cuda") targets[:, :128, :128] = torch.randint(0, 21, (8, 128, 128), device="cuda") loss = loss_fn(logits, targets)

Initialise CombinedSegmentationLoss with term weights and ignore index.

Source code in src/corecv/losses/segmentation.py

def __init__(
    self,
    ce_weight: float = 1.0,
    dice_weight: float = 1.0,
    ignore_index: int = -100,
) -> None:
    """Initialise CombinedSegmentationLoss with term weights and ignore index."""
    super().__init__()
    self.ce_weight = float(ce_weight)
    self.dice_weight = float(dice_weight)
    self.ignore_index = int(ignore_index)
    self._dice_fn = DiceLoss(smooth=1.0, reduction="mean")

forward(inputs, targets)

Compute the combined CE + Dice loss.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Logits of shape (B, C, H, W).	required
targets	Tensor	Ground-truth class indices of shape (B, H, W). Pixels equal to self.ignore_index are excluded from both loss terms.	required

Returns:

Type	Description
Tensor	Weighted scalar loss.

Source code in src/corecv/losses/segmentation.py

def forward(self, inputs: Tensor, targets: Tensor) -> Tensor:
    """Compute the combined CE + Dice loss.

    Args:
        inputs: Logits of shape ``(B, C, H, W)``.
        targets: Ground-truth class indices of shape ``(B, H, W)``.
            Pixels equal to ``self.ignore_index`` are excluded from
            both loss terms.

    Returns:
        Weighted scalar loss.
    """
    # ---- Cross-entropy component ------------------------------------
    ce_loss: Tensor = F.cross_entropy(
        inputs,
        targets,
        ignore_index=self.ignore_index,
        reduction="mean",
    )

    # ---- Dice component with ignore-mask ----------------------------
    # Build a boolean mask of valid (non-ignored) pixels.
    valid_mask: Tensor = targets != self.ignore_index  # (B, H, W)

    # Replace ignored targets with a safe dummy class (0) so that
    # F.one_hot does not produce out-of-range indices.  The mask
    # ensures these positions contribute zero to the Dice sum.
    safe_targets: Tensor = targets.clone()
    safe_targets[~valid_mask] = 0

    num_classes: int = inputs.shape[1]

    # Softmax predictions
    probs: Tensor = F.softmax(inputs, dim=1)

    # One-hot targets: (B, H, W) -> (B, H, W, C) -> (B, C, H, W)
    one_hot: Tensor = F.one_hot(safe_targets, num_classes=num_classes)
    one_hot = one_hot.permute(0, 3, 1, 2).float()

    # Apply the valid mask to both predictions and targets so that
    # ignored positions contribute exactly zero to all sums.
    # valid_mask: (B, H, W) -> (B, 1, H, W) for broadcasting
    mask_4d: Tensor = valid_mask.unsqueeze(1).float()  # (B, 1, H, W)
    probs_masked: Tensor = probs * mask_4d
    target_masked: Tensor = one_hot * mask_4d

    # Vectorised Dice over valid pixels only
    probs_flat: Tensor = probs_masked.flatten(2)   # (B, C, N)
    target_flat: Tensor = target_masked.flatten(2)  # (B, C, N)

    intersection: Tensor = (probs_flat * target_flat).sum(dim=2)
    probs_sum: Tensor = probs_flat.sum(dim=2)
    target_sum: Tensor = target_flat.sum(dim=2)

    smooth = 1.0
    dice: Tensor = (2.0 * intersection + smooth) / (
        probs_sum + target_sum + smooth
    )
    dice_loss: Tensor = (1.0 - dice).mean(dim=1)  # per-sample, mean over classes

    # If all pixels in a sample are ignored, the dice loss is 0
    # (division by zero is avoided because smooth > 0).
    dice_loss_mean: Tensor = dice_loss.mean()

    return self.ce_weight * ce_loss + self.dice_weight * dice_loss_mean

DiceLoss(smooth=1.0, reduction='mean')

Bases: Module

Dice coefficient loss for semantic segmentation.

The Dice coefficient measures overlap between prediction and target::

Dice = 2 * sum(p * t) / (sum(p) + sum(t) + smooth)

The loss is 1 - Dice. A value of smooth (default 1.0) prevents division by zero.

This implementation is fully vectorised: F.one_hot expands targets to (B, H, W, C) in a single kernel call, then the Dice computation is a sequence of batch reductions with no Python loops over channels or spatial locations.

Parameters:

Name	Type	Description	Default
smooth	float	Laplace smoothing factor to avoid division by zero. Default 1.0.	1.0
reduction	str	Aggregation mode ('mean' | 'sum' | 'none'). Default 'mean'.	'mean'

Note

Inputs should be logits (not probabilities). The module applies channel-wise softmax internally.

Example

loss_fn = DiceLoss(smooth=1.0) logits = torch.randn(2, 10, 128, 128, device="cuda") targets = torch.randint(0, 10, (2, 128, 128), device="cuda") loss = loss_fn(logits, targets)

Initialise DiceLoss with smoothing factor and reduction mode.

Source code in src/corecv/losses/segmentation.py

def __init__(
    self,
    smooth: float = 1.0,
    reduction: str = "mean",
) -> None:
    """Initialise DiceLoss with smoothing factor and reduction mode."""
    super().__init__()
    if reduction not in ("mean", "sum", "none"):
        msg = f"Invalid reduction: {reduction!r}. Must be 'mean', 'sum', or 'none'."
        raise ValueError(msg)

    self.smooth = float(smooth)
    self.reduction = reduction

forward(inputs, targets)

Compute Dice loss.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Logits of shape (B, C, H, W) where C is the number of segmentation classes.	required
targets	Tensor	Ground-truth class indices of shape (B, H, W) with values in [0, C).	required

Returns:

Type	Description
Tensor	Scalar or tensor loss depending on self.reduction.

Source code in src/corecv/losses/segmentation.py

def forward(self, inputs: Tensor, targets: Tensor) -> Tensor:
    """Compute Dice loss.

    Args:
        inputs: Logits of shape ``(B, C, H, W)`` where ``C`` is the
            number of segmentation classes.
        targets: Ground-truth class indices of shape ``(B, H, W)`` with
            values in ``[0, C)``.

    Returns:
        Scalar or tensor loss depending on ``self.reduction``.
    """
    num_classes: int = inputs.shape[1]

    # --- Softmax predictions: (B, C, H, W) -> probabilities --------
    probs: Tensor = F.softmax(inputs, dim=1)

    # --- One-hot encode targets: (B, H, W) -> (B, H, W, C) --------
    # F.one_hot expects Long dtype
    targets_long = targets.long()
    one_hot: Tensor = F.one_hot(targets_long, num_classes=num_classes)
    # (B, H, W, C) -> (B, C, H, W) to match probs layout
    one_hot = one_hot.permute(0, 3, 1, 2).float()

    # --- Vectorised Dice computation --------------------------------
    # Flatten spatial dims: (B, C, H, W) -> (B, C, H*W)
    probs_flat: Tensor = probs.flatten(2)      # (B, C, N)
    target_flat: Tensor = one_hot.flatten(2)    # (B, C, N)

    # Intersection: sum of element-wise product across spatial dim
    # (B, C, N) -> (B, C)  [sum over N]
    intersection: Tensor = (probs_flat * target_flat).sum(dim=2)

    probs_sum: Tensor = probs_flat.sum(dim=2)
    target_sum: Tensor = target_flat.sum(dim=2)

    # Dice coefficient per (batch, class): (B, C)
    dice: Tensor = (2.0 * intersection + self.smooth) / (
        probs_sum + target_sum + self.smooth
    )

    # Dice loss per (batch, class): (B, C)
    dice_loss: Tensor = 1.0 - dice

    # --- Per-sample reduction across classes ------------------------
    # Mean over classes to get per-sample loss: (B,)
    per_sample_loss: Tensor = dice_loss.mean(dim=1)

    return self._reduce(per_sample_loss)

---

Assigners

Hungarian Assigner

corecv.losses.assigners.hungarian

GPU-native Hungarian matching and set criterion for query detection.

Implements bipartite 1-to-1 assignment (Hungarian algorithm) and the associated SetCriterion used to train RT-DETR / D-FINE style query-based detection heads.

All loss computations (focal classification, L1 bounding-box, GIoU) are pure vectorised PyTorch operations with zero CPU-GPU synchronisations during forward and backward passes. The only CPU round-trip is the combinatorial assignment step itself, which operates on a small cost matrix of shape (num_queries, num_gt) per batch element — the heavy lifting (the full pairwise cost-matrix construction) stays entirely on VRAM.

Architecture overview::

pred_scores:  (B, num_queries, num_classes)
pred_boxes:   (B, num_queries, 4)              — normalised cxcywh
gt_labels:    List[Tensor]  of per-image targets
gt_boxes:     List[Tensor]  of per-image targets

    |
    v  HungarianMatcher.forward()  (cost_class + L1 + GIoU)
    |
indices: List[Tuple[Tensor, Tensor]]  (pred_idx, gt_idx) per image

    |
    v  SetCriterion.loss()  (focal + L1 + GIoU + aux)
    |
scalar loss

Example

import torch from corecv.losses.assigners import HungarianMatcher, SetCriterion matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0) criterion = SetCriterion(num_classes=80, matcher=matcher) pred_scores = torch.randn(2, 300, 80) pred_boxes = torch.sigmoid(torch.randn(2, 300, 4)) gt_labels = [torch.tensor([0, 5]), torch.tensor([1, 2, 3])] gt_boxes = [torch.rand(2, 4), torch.rand(3, 4)] loss, aux = criterion(pred_scores, pred_boxes, gt_labels, gt_boxes)

HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, focal_alpha=0.25, focal_gamma=2.0)

Bases: Module

Bipartite 1-to-1 assignment via the Hungarian algorithm.

Computes a cost matrix from three terms — classification cost, L1 bounding-box cost, and GIoU cost — then solves the linear-sum assignment to find the optimal prediction-to-ground-truth pairing.

The cost matrix for each image is::

cost = cost_class * focal_cost
     + cost_bbox * l1_cost
     + cost_giou * (1 - giou_matrix)

where:

• focal_cost is derived from the negative focal-log-probability for the target class.
• l1_cost is the element-wise L1 distance between normalised (cx, cy, w, h) boxes.
• giou_matrix is the pairwise GIoU between all predictions and all ground-truth boxes.

The assignment is performed per image independently. The heavy pairwise-cost computation runs entirely on VRAM; only the small (num_queries, num_gt) cost matrix is transferred to CPU for scipy.optimize.linear_sum_assignment.

Parameters:

Name	Type	Description	Default
cost_class	float	Weight of the focal classification cost. Default 2.0.	2.0
cost_bbox	float	Weight of the L1 bounding-box cost. Default 5.0.	5.0
cost_giou	float	Weight of the GIoU cost. Default 2.0.	2.0
focal_alpha	float	Focal-loss alpha for class weighting. Default 0.25.	0.25
focal_gamma	float	Focal-loss gamma for focusing parameter. Default 2.0.	2.0

Example

matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0) pred_scores = torch.randn(2, 300, 80) pred_boxes = torch.sigmoid(torch.randn(2, 300, 4)) gt_labels = [torch.tensor([0, 5]), torch.tensor([1, 2, 3])] gt_boxes = [torch.rand(2, 4), torch.rand(3, 4)] indices = matcher(pred_scores, pred_boxes, gt_labels, gt_boxes)

Initialise HungarianMatcher with cost weights and focal parameters.

Source code in src/corecv/losses/assigners/hungarian.py

def __init__(  # noqa: PLR0913
    self,
    cost_class: float = 2.0,
    cost_bbox: float = 5.0,
    cost_giou: float = 2.0,
    focal_alpha: float = 0.25,
    focal_gamma: float = 2.0,
) -> None:
    """Initialise HungarianMatcher with cost weights and focal parameters."""
    super().__init__()
    self.cost_class = float(cost_class)
    self.cost_bbox = float(cost_bbox)
    self.cost_giou = float(cost_giou)
    self.focal_alpha = float(focal_alpha)
    self.focal_gamma = float(focal_gamma)

forward(pred_scores, pred_boxes, gt_labels, gt_boxes)

Compute optimal bipartite matching for each image in the batch.

Parameters:

Name	Type	Description	Default
pred_scores	Tensor	Classification logits of shape (B, num_queries, num_classes).	required
pred_boxes	Tensor	Predicted normalised (cx, cy, w, h) boxes of shape (B, num_queries, 4).	required
gt_labels	list[Tensor]	List of B 1-D integer tensors, each containing the class labels for that image's ground-truth boxes.	required
gt_boxes	list[Tensor]	List of B 2-D tensors of shape (num_gt_i, 4) with normalised (cx, cy, w, h) ground-truth boxes.	required

Returns:

Type	Description
list[tuple[Tensor, Tensor]]	List of B tuples (pred_indices, gt_indices) where each
list[tuple[Tensor, Tensor]]	element is a 1-D LongTensor of matched indices. If an
list[tuple[Tensor, Tensor]]	image has zero ground-truth boxes, both tensors are empty.

Source code in src/corecv/losses/assigners/hungarian.py

def forward(
    self,
    pred_scores: Tensor,
    pred_boxes: Tensor,
    gt_labels: list[Tensor],
    gt_boxes: list[Tensor],
) -> list[tuple[Tensor, Tensor]]:
    """Compute optimal bipartite matching for each image in the batch.

    Args:
        pred_scores: Classification logits of shape
            ``(B, num_queries, num_classes)``.
        pred_boxes: Predicted normalised ``(cx, cy, w, h)`` boxes of
            shape ``(B, num_queries, 4)``.
        gt_labels: List of ``B`` 1-D integer tensors, each containing
            the class labels for that image's ground-truth boxes.
        gt_boxes: List of ``B`` 2-D tensors of shape ``(num_gt_i, 4)``
            with normalised ``(cx, cy, w, h)`` ground-truth boxes.

    Returns:
        List of ``B`` tuples ``(pred_indices, gt_indices)`` where each
        element is a 1-D ``LongTensor`` of matched indices.  If an
        image has zero ground-truth boxes, both tensors are empty.
    """
    B = pred_scores.shape[0]
    device = pred_scores.device
    indices: list[tuple[Tensor, Tensor]] = []

    for i in range(B):
        # Per-image tensors — still on GPU, no sync yet
        p_scores_i = pred_scores[i]  # (num_queries, C)
        p_boxes_i = pred_boxes[i]    # (num_queries, 4)
        gt_label_i = gt_labels[i]    # (num_gt,)
        gt_box_i = gt_boxes[i]       # (num_gt, 4)

        num_gt = gt_box_i.shape[0]
        if num_gt == 0:
            # No ground-truth — empty assignment
            indices.append(
                (
                    torch.zeros(0, dtype=torch.long, device=device),
                    torch.zeros(0, dtype=torch.long, device=device),
                )
            )
            continue

        # ---- Pairwise cost matrix (fully GPU-native) ---------------
        cost_matrix = self._compute_cost_matrix(
            p_scores_i, p_boxes_i, gt_label_i, gt_box_i,
        )  # (num_queries, num_gt)

        # ---- Solve assignment on CPU (small matrix) -----------------
        # Transfer only the cost matrix — all heavy computation was
        # performed on GPU.
        cost_np = cost_matrix.detach().cpu().numpy()
        row_ind, col_ind = linear_sum_assignment(cost_np)

        indices.append(
            (
                torch.as_tensor(row_ind, dtype=torch.long, device=device),
                torch.as_tensor(col_ind, dtype=torch.long, device=device),
            )
        )

    return indices

SetCriterion(num_classes, matcher, focal_alpha=0.25, focal_gamma=2.0, loss_cls_weight=1.0, loss_bbox_weight=5.0, loss_giou_weight=2.0, aux_loss_weights=None, no_object_weight=0.1)

Bases: Module

Set-based criterion for training query detection heads (RT-DETR / D-FINE).

Uses :class:HungarianMatcher to establish 1-to-1 assignments between predictions and ground-truth boxes, then computes a weighted sum of:

• Classification loss — sigmoid focal loss applied to matched queries; unmatched queries are treated as background (class 0 with target 0 in the "no-object" formulation, or omitted entirely).
• Bounding-box L1 loss — element-wise L1 on matched boxes only.
• GIoU loss — 1 - GIoU on matched boxes only.

Supports auxiliary losses from intermediate decoder layers: when return_intermediate=True is used in the head, pass the lists of intermediate predictions and the criterion will apply the same loss at every decoder layer with optional per-layer weights.

All loss computations are 100 % GPU-native — the only CPU round-trip is the assignment step inside :class:HungarianMatcher.

Parameters:

Name	Type	Description	Default
num_classes	int	Number of foreground object classes (excluding the implicit "no-object" class).	required
matcher	HungarianMatcher	A :class:HungarianMatcher instance.	required
focal_alpha	float	Focal-loss alpha. Default 0.25.	0.25
focal_gamma	float	Focal-loss gamma. Default 2.0.	2.0
loss_cls_weight	float	Weight of the classification loss. Default 1.0.	1.0
loss_bbox_weight	float	Weight of the L1 bounding-box loss. Default 5.0.	5.0
loss_giou_weight	float	Weight of the GIoU loss. Default 2.0.	2.0
aux_loss_weights	list[float] | None	Per-layer weights for auxiliary decoder losses. If None, all layers are weighted equally. The length must match the number of intermediate predictions when auxiliary losses are used.	None
no_object_weight	float	Weight assigned to the "no-object" class for unmatched queries. Default 0.1.	0.1

Example

matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0) criterion = SetCriterion(num_classes=80, matcher=matcher) pred_scores = torch.randn(2, 300, 80) pred_boxes = torch.sigmoid(torch.randn(2, 300, 4)) gt_labels = [torch.tensor([0, 5]), torch.tensor([1, 2, 3])] gt_boxes = [torch.rand(2, 4), torch.rand(3, 4)] loss, aux = criterion(pred_scores, pred_boxes, gt_labels, gt_boxes)

Initialise SetCriterion with matcher and loss weights.

Source code in src/corecv/losses/assigners/hungarian.py

def __init__(  # noqa: PLR0913
    self,
    num_classes: int,
    matcher: HungarianMatcher,
    focal_alpha: float = 0.25,
    focal_gamma: float = 2.0,
    loss_cls_weight: float = 1.0,
    loss_bbox_weight: float = 5.0,
    loss_giou_weight: float = 2.0,
    aux_loss_weights: list[float] | None = None,
    no_object_weight: float = 0.1,
) -> None:
    """Initialise SetCriterion with matcher and loss weights."""
    super().__init__()
    self.num_classes = int(num_classes)
    self.matcher = matcher
    self.focal_alpha = float(focal_alpha)
    self.focal_gamma = float(focal_gamma)
    self.loss_cls_weight = float(loss_cls_weight)
    self.loss_bbox_weight = float(loss_bbox_weight)
    self.loss_giou_weight = float(loss_giou_weight)
    self.aux_loss_weights = aux_loss_weights
    self.no_object_weight = float(no_object_weight)

forward(pred_scores, pred_boxes, gt_labels, gt_boxes, intermediate_cls=None, intermediate_reg=None)

Compute the set-matching loss.

Parameters:

Name	Type	Description	Default
pred_scores	Tensor	Final classification logits (B, num_queries, num_classes).	required
pred_boxes	Tensor	Final predicted normalised (cx, cy, w, h) boxes of shape (B, num_queries, 4).	required
gt_labels	list[Tensor]	List of B 1-D integer tensors with per-image ground-truth class labels.	required
gt_boxes	list[Tensor]	List of B 2-D tensors of shape (num_gt_i, 4) with normalised ground-truth boxes.	required
intermediate_cls	list[Tensor] | None	Optional list of intermediate classification logits from decoder layers, each of shape (B, num_queries, num_classes).	None
intermediate_reg	list[Tensor] | None	Optional list of intermediate bounding-box predictions from decoder layers, each of shape (B, num_queries, 4).	None

Returns:

Type	Description
Tensor	Tuple of:
dict[str, Tensor]	loss: Scalar combined loss tensor.
tuple[Tensor, dict[str, Tensor]]	aux_losses: Dictionary mapping "loss_cls", "loss_bbox", "loss_giou" to their scalar values. Auxiliary layer losses are stored under "aux_loss_cls", "aux_loss_bbox", "aux_loss_giou".

Source code in src/corecv/losses/assigners/hungarian.py

def forward(  # noqa: PLR0913
    self,
    pred_scores: Tensor,
    pred_boxes: Tensor,
    gt_labels: list[Tensor],
    gt_boxes: list[Tensor],
    intermediate_cls: list[Tensor] | None = None,
    intermediate_reg: list[Tensor] | None = None,
) -> tuple[Tensor, dict[str, Tensor]]:
    """Compute the set-matching loss.

    Args:
        pred_scores: Final classification logits
            ``(B, num_queries, num_classes)``.
        pred_boxes: Final predicted normalised ``(cx, cy, w, h)``
            boxes of shape ``(B, num_queries, 4)``.
        gt_labels: List of ``B`` 1-D integer tensors with per-image
            ground-truth class labels.
        gt_boxes: List of ``B`` 2-D tensors of shape
            ``(num_gt_i, 4)`` with normalised ground-truth boxes.
        intermediate_cls: Optional list of intermediate classification
            logits from decoder layers, each of shape
            ``(B, num_queries, num_classes)``.
        intermediate_reg: Optional list of intermediate bounding-box
            predictions from decoder layers, each of shape
            ``(B, num_queries, 4)``.

    Returns:
        Tuple of:

        * ``loss``: Scalar combined loss tensor.
        * ``aux_losses``: Dictionary mapping ``"loss_cls"``,
            ``"loss_bbox"``, ``"loss_giou"`` to their scalar values.
            Auxiliary layer losses are stored under
            ``"aux_loss_cls"``, ``"aux_loss_bbox"``, ``"aux_loss_giou"``.
    """
    # ---- Run Hungarian matching on the final predictions ------------
    indices = self.matcher(pred_scores, pred_boxes, gt_labels, gt_boxes)

    # ---- Compute losses for the final decoder layer ----------------
    loss_cls, loss_bbox, loss_giou = self._compute_losses(
        pred_scores, pred_boxes, indices, gt_labels, gt_boxes,
    )

    total_cls = self.loss_cls_weight * loss_cls
    total_bbox = self.loss_bbox_weight * loss_bbox
    total_giou = self.loss_giou_weight * loss_giou
    total_loss = total_cls + total_bbox + total_giou

    aux_losses: dict[str, Tensor] = {
        "loss_cls": loss_cls.detach(),
        "loss_bbox": loss_bbox.detach(),
        "loss_giou": loss_giou.detach(),
    }

    # ---- Auxiliary losses from intermediate decoder layers ----------
    if intermediate_cls is not None and intermediate_reg is not None:
        num_layers = len(intermediate_cls)
        weights = self.aux_loss_weights
        if weights is None:
            # Equal weighting for all auxiliary layers
            weights = [1.0 / num_layers] * num_layers

        if len(weights) != num_layers:
            msg = (
                f"aux_loss_weights length ({len(weights)}) must match "
                f"number of intermediate layers ({num_layers})."
            )
            raise ValueError(msg)

        for layer_idx, (cls_i, reg_i) in enumerate(
            zip(intermediate_cls, intermediate_reg, strict=True),
        ):
            # Re-run matching at each intermediate layer
            aux_indices = self.matcher(cls_i, reg_i, gt_labels, gt_boxes)
            aux_cls, aux_bbox, aux_giou = self._compute_losses(
                cls_i, reg_i, aux_indices, gt_labels, gt_boxes,
            )
            w = weights[layer_idx]
            total_loss = total_loss + w * (
                self.loss_cls_weight * aux_cls
                + self.loss_bbox_weight * aux_bbox
                + self.loss_giou_weight * aux_giou
            )
            aux_losses[f"aux_loss_cls_layer{layer_idx}"] = aux_cls.detach()
            aux_losses[f"aux_loss_bbox_layer{layer_idx}"] = aux_bbox.detach()
            aux_losses[f"aux_loss_giou_layer{layer_idx}"] = aux_giou.detach()

    return total_loss, aux_losses

Task-Aligned Assigner (TAL)

corecv.losses.assigners.tal

GPU-native Task-Aligned Assigner (TAL) for anchor-free detection.

Implements the dynamic sample assignment strategy from TOOD: Task-Aligned One-stage Object Detection (Tian et al., AAAI 2021). The alignment metric used for positive sample selection is::

t = cls_score^alpha * IoU^beta

For each ground-truth box, the top-k predictions with the highest TAL alignment metric are selected as positive samples. All remaining predictions are treated as negatives.

The assigner is designed to integrate with :class:~corecv.models.heads.detection.decoupled_anchor_free.DecoupledAnchorFreeHead and downstream losses (:class:~corecv.losses.detection.QualityFocalLoss / :class:~corecv.losses.detection.VarifocalLoss for classification, :class:~corecv.losses.detection.GIoULoss / :class:~corecv.losses.detection.CIoULoss for bounding-box regression).

All computations are pure vectorised PyTorch operations with zero CPU-GPU synchronisations during assignment. No .item() calls, no .cpu() transfers — everything stays on VRAM.

Architecture overview::

pred_scores:  List[(B, C, H_l, W_l)]  — per-level classification logits
pred_boxes:   List[(B, 4, H_l, W_l)]  — per-level (l, t, r, b) regression
strides:      List[int]               — per-level strides (e.g. 8, 16, 32)
gt_labels:    List[Tensor]            — per-image GT class labels
gt_boxes:     List[Tensor]            — per-image GT boxes (x1, y1, x2, y2)

    |
    v  TaskAlignedAssigner.forward()
    |
assignment dict with:
    pos_mask       — (N_total,) bool tensor per image
    neg_mask       — (N_total,) bool tensor per image
    assigned_gt_inds — (N_total,) long tensor per image
    assigned_labels  — (N_total,) long tensor per image
    pos_ious       — (num_pos,) float tensor per image (IoU quality)

Example

from corecv.losses.assigners.tal import TaskAlignedAssigner assigner = TaskAlignedAssigner(num_classes=80, topk=13)

Given per-level predictions from DecoupledAnchorFreeHead:

cls_logits: [(B, 80, H_l, W_l), ...], reg_pred: [(B, 4, H_l, W_l), ...]

strides: [8, 16, 32], gt_labels: per-image, gt_boxes: per-image (xyxy)

assignment = assigner(cls_logits, reg_pred, strides, gt_labels, gt_boxes)

TaskAlignedAssigner(num_classes, topk=13, alpha=0.5, beta=6.0)

Bases: Module

Task-Aligned dynamic assigner for anchor-free detection.

Implements the alignment-metric-based sample assignment from TOOD. For each ground-truth box, the topk predictions with the highest alignment score t = cls^alpha * IoU^beta are selected as positives. When a prediction is selected by multiple ground-truth boxes, it is assigned to the one yielding the highest alignment metric.

The assigner produces per-image assignment results that can be consumed directly by the CoreCV detection losses:

• :class:~corecv.losses.detection.QualityFocalLoss or :class:~corecv.losses.detection.VarifocalLoss for classification (using assigned_labels and pos_ious as quality targets).
• :class:~corecv.losses.detection.GIoULoss or :class:~corecv.losses.detection.CIoULoss for bounding-box regression (using pos_mask and assigned_gt_inds to gather matched predictions and targets).

All operations run entirely on GPU — there are zero CPU-GPU synchronisations during the forward pass.

Parameters:

Name	Type	Description	Default
num_classes	int	Number of foreground object classes.	required
topk	int	Number of top predictions selected per ground-truth box. Typical values: 9 for stride-8 only, 13 for multi-scale. Default 13.	13
alpha	float	Exponent for the classification score in the alignment metric. Default 0.5.	0.5
beta	float	Exponent for the IoU in the alignment metric. Higher values bias selection towards higher-quality boxes. Default 6.0.	6.0

Raises:

Type	Description
ValueError	If topk is less than 1, num_classes is less than 1, or alpha / beta are negative.

Example

assigner = TaskAlignedAssigner(num_classes=80, topk=13) cls_logits = [torch.randn(2, 80, 80, 80), ... torch.randn(2, 80, 40, 40), ... torch.randn(2, 80, 20, 20)] reg_pred = [torch.randn(2, 4, 80, 80).abs(), ... torch.randn(2, 4, 40, 40).abs(), ... torch.randn(2, 4, 20, 20).abs()] strides = [8, 16, 32] gt_labels = [torch.tensor([0, 5]), torch.tensor([1, 2, 3])] gt_boxes = [torch.rand(2, 4) * 400 + 50, ... torch.rand(3, 4) * 400 + 50] assignment = assigner(cls_logits, reg_pred, strides, ... gt_labels, gt_boxes)

Initialise the Task-Aligned Assigner.

Source code in src/corecv/losses/assigners/tal.py

def __init__(  # noqa: PLR0913
    self,
    num_classes: int,
    topk: int = 13,
    alpha: float = 0.5,
    beta: float = 6.0,
) -> None:
    """Initialise the Task-Aligned Assigner."""
    super().__init__()
    if num_classes < 1:
        msg = f"num_classes must be >= 1, got {num_classes}."
        raise ValueError(msg)
    if topk < _MIN_TOPK:
        msg = f"topk must be >= {_MIN_TOPK}, got {topk}."
        raise ValueError(msg)
    if alpha < 0.0:
        msg = f"alpha must be >= 0, got {alpha}."
        raise ValueError(msg)
    if beta < 0.0:
        msg = f"beta must be >= 0, got {beta}."
        raise ValueError(msg)

    self.num_classes = int(num_classes)
    self.topk = int(topk)
    self.alpha = float(alpha)
    self.beta = float(beta)

forward(pred_scores, pred_boxes, strides, gt_labels, gt_boxes)

Compute task-aligned assignment for each image in the batch.

Parameters:

Name	Type	Description	Default
pred_scores	list[Tensor]	Per-level classification logits. Each element is a tensor of shape (B, C, H_l, W_l) where C equals self.num_classes.	required
pred_boxes	list[Tensor]	Per-level bounding-box regression predictions. Each element is a tensor of shape (B, 4, H_l, W_l) in (l, t, r, b) format (distances from cell centres).	required
strides	list[int]	Per-level stride values (e.g. [8, 16, 32]). len(strides) must equal len(pred_scores).	required
gt_labels	list[Tensor]	List of B 1-D integer tensors, each containing the class labels for that image's ground-truth boxes. Values in [0, num_classes).	required
gt_boxes	list[Tensor]	List of B 2-D tensors of shape (num_gt_i, 4) with ground-truth boxes in absolute (x1, y1, x2, y2) pixel coordinates.	required

Returns:

Type	Description
dict[str, list[Tensor]]	Dictionary with the following keys (all values are lists
dict[str, list[Tensor]]	of length B, one element per image):
dict[str, list[Tensor]]	"pos_mask": bool tensor of shape (N_i,) — True for positive predictions.
dict[str, list[Tensor]]	"neg_mask": bool tensor of shape (N_i,) — True for negative predictions.
dict[str, list[Tensor]]	"assigned_gt_inds": int64 tensor of shape (N_i,) — ground-truth index assigned to each prediction (-1 for negatives).
dict[str, list[Tensor]]	"assigned_labels": int64 tensor of shape (N_i,) — class label assigned to each prediction (0 for negatives).
dict[str, list[Tensor]]	"pos_ious": float32 tensor of shape (num_pos_i,) — IoU quality of each positive prediction with its assigned ground-truth box.
dict[str, list[Tensor]]	Here N_i = sum(H_l * W_l) across all levels for image
dict[str, list[Tensor]]	i.

Source code in src/corecv/losses/assigners/tal.py

def forward(
    self,
    pred_scores: list[Tensor],
    pred_boxes: list[Tensor],
    strides: list[int],
    gt_labels: list[Tensor],
    gt_boxes: list[Tensor],
) -> dict[str, list[Tensor]]:
    """Compute task-aligned assignment for each image in the batch.

    Args:
        pred_scores: Per-level classification logits.  Each element
            is a tensor of shape ``(B, C, H_l, W_l)`` where ``C``
            equals ``self.num_classes``.
        pred_boxes: Per-level bounding-box regression predictions.
            Each element is a tensor of shape ``(B, 4, H_l, W_l)``
            in ``(l, t, r, b)`` format (distances from cell centres).
        strides: Per-level stride values (e.g. ``[8, 16, 32]``).
            ``len(strides)`` must equal ``len(pred_scores)``.
        gt_labels: List of ``B`` 1-D integer tensors, each
            containing the class labels for that image's
            ground-truth boxes.  Values in ``[0, num_classes)``.
        gt_boxes: List of ``B`` 2-D tensors of shape
            ``(num_gt_i, 4)`` with ground-truth boxes in absolute
            ``(x1, y1, x2, y2)`` pixel coordinates.

    Returns:
        Dictionary with the following keys (all values are lists
        of length ``B``, one element per image):

        * ``"pos_mask"``: ``bool`` tensor of shape ``(N_i,)`` —
            ``True`` for positive predictions.
        * ``"neg_mask"``: ``bool`` tensor of shape ``(N_i,)`` —
            ``True`` for negative predictions.
        * ``"assigned_gt_inds"``: ``int64`` tensor of shape
            ``(N_i,)`` — ground-truth index assigned to each
            prediction (``-1`` for negatives).
        * ``"assigned_labels"``: ``int64`` tensor of shape
            ``(N_i,)`` — class label assigned to each prediction
            (``0`` for negatives).
        * ``"pos_ious"``: ``float32`` tensor of shape
            ``(num_pos_i,)`` — IoU quality of each positive
            prediction with its assigned ground-truth box.

        Here ``N_i = sum(H_l * W_l)`` across all levels for image
        ``i``.
    """
    B = pred_scores[0].shape[0]
    device = pred_scores[0].device
    num_levels = len(pred_scores)

    # Validate input consistency
    if len(pred_boxes) != num_levels:
        msg = (
            f"pred_scores and pred_boxes must have the same length, "
            f"got {num_levels} and {len(pred_boxes)}."
        )
        raise ValueError(msg)
    if len(strides) != num_levels:
        msg = (
            f"strides length ({len(strides)}) must match the number "
            f"of feature levels ({num_levels})."
        )
        raise ValueError(msg)

    # Compute feature map spatial sizes from the regression outputs
    feat_sizes: list[tuple[int, int]] = [
        (reg.shape[_DIM_4D - 2], reg.shape[_DIM_4D - 1])  # (H, W)
        for reg in pred_boxes
    ]

    # Generate anchor centres for all levels — (N_total, 2)
    all_anchors: Tensor = _make_anchors(strides, feat_sizes, device)

    # Flatten predictions across all levels for each image
    cls_scores: Tensor = torch.cat(
        [
            level_cls.permute(0, 2, 3, 1).reshape(B, -1, self.num_classes)
            for level_cls in pred_scores
        ],
        dim=1,
    )
    reg_preds: Tensor = torch.cat(
        [
            level_reg.permute(0, 2, 3, 1).reshape(B, -1, 4)
            for level_reg in pred_boxes
        ],
        dim=1,
    )

    # Decode all predictions from (l, t, r, b) to absolute (x1, y1, x2, y2)
    decoded_boxes: Tensor = _decode_boxes(
        reg_preds.reshape(-1, 4),  # (B*N_total, 4)
        all_anchors.unsqueeze(0).expand(B, -1, -1).reshape(-1, 2),  # (B*N_total, 2)
    ).reshape(B, all_anchors.shape[0], 4)  # (B, N_total, 4)

    # Apply sigmoid to get classification probabilities
    cls_probs: Tensor = cls_scores.sigmoid()  # (B, N_total, C)

    # Per-image assignment
    results: dict[str, list[Tensor]] = {
        "pos_mask": [],
        "neg_mask": [],
        "assigned_gt_inds": [],
        "assigned_labels": [],
        "pos_ious": [],
    }

    for i in range(B):
        self._assign_single_image(
            cls_probs_i=cls_probs[i],
            decoded_boxes_i=decoded_boxes[i],
            gt_labels_i=gt_labels[i],
            gt_boxes_i=gt_boxes[i],
            results=results,
            device=device,
        )

    return results