Class weight for CrossEntropy loss in Faster RCNN

🚀 The feature

Hi, in order to give more importance to some classes, add a parameter to implement class weight in torchvision.models.detection.roi_heads.fastrcnn_loss.

Motivation, pitch

It would be useful in unbalanced datasets

Alternatives

I propose something like that:

def fastrcnn_loss(class_logits, box_regression, labels, regression_targets, pos_weight):
    # type: (Tensor, Tensor, List[Tensor], List[Tensor], pos_weight) -> Tuple[Tensor, Tensor]
    """
    Computes the loss for Faster R-CNN.
    Args:
        class_logits (Tensor)
        box_regression (Tensor)
        labels (list[BoxList])
        regression_targets (Tensor)
        pos_weight (Tensor)

    Returns:
        classification_loss (Tensor)
        box_loss (Tensor)
    """

    labels = torch.cat(labels, dim=0)
    regression_targets = torch.cat(regression_targets, dim=0)
    pos_weight = pos_weight.cuda() if labels.is_cuda else pos_weight
    classification_loss = F.cross_entropy(class_logits, labels, weight=pos_weight)

    # get indices that correspond to the regression targets for
    # the corresponding ground truth labels, to be used with
    # advanced indexing
    sampled_pos_inds_subset = torch.where(labels > 0)[0]
    labels_pos = labels[sampled_pos_inds_subset]
    N, num_classes = class_logits.shape
    box_regression = box_regression.reshape(N, box_regression.size(-1) // 4, 4)

    box_loss = F.smooth_l1_loss(
        box_regression[sampled_pos_inds_subset, labels_pos],
        regression_targets[sampled_pos_inds_subset],
        beta=1 / 9,
        reduction="sum",
    )
    box_loss = box_loss / labels.numel()

    return classification_loss, 

Additional context

No response

Start in torchvision/models/detection/roi_heads.py at fastrcnn_loss and trace its callers to determine how a class-weight parameter would be supplied. Review the existing CrossEntropy loss path and the proposed behavior in the issue. Done means Faster R-CNN can apply class weights for classification loss without changing box regression behavior.