Generalized Focal Loss: Learning Qualified and Distributed Bounding Boxes for Dense Object Detection

def quality_focal_loss(pred, target, beta=2.0):
    """
    Args:
        pred (torch.Tensor): Predicted joint representation of classification
            and quality (IoU) estimation with shape (N, C), C is the number of
            classes.
        target (tuple([torch.Tensor])): Target category label with shape (N,)
            and target quality label with shape (N,).
        beta (float): The beta parameter for calculating the modulating factor.
            Defaults to 2.0.
    Returns:
        torch.Tensor: Loss tensor with shape (N,).
    """
    assert len(target) == 2, """target for QFL must be a tuple of two elements,
        including category label and quality label, respectively"""
    # label denotes the category id, score denotes the quality score
    label, score = target  # label(N) 类别,score(N),预测bbox和gt bbox的iou

    # negatives are supervised by 0 quality score
    pred_sigmoid = pred.sigmoid()  # (N,class_num)
    scale_factor = pred_sigmoid
    zerolabel = scale_factor.new_zeros(pred.shape)
    # 先假设所有label都是负样本，计算bce loss，乘上sigmoid^beta次方，达到focal效应
    loss = F.binary_cross_entropy_with_logits(
        pred, zerolabel, reduction='none') * scale_factor.pow(beta)

    # FG cat_id: [0, num_classes -1], BG cat_id: num_classes
    bg_class_ind = pred.size(1)
    # 找出正样本索引
    pos = ((label >= 0) & (label < bg_class_ind)).nonzero().squeeze(1)
    pos_label = label[pos].long()
    # positives are supervised by bbox quality (IoU) score
    scale_factor = score[pos] - pred_sigmoid[pos, pos_label]  # 正样本focal值
    # 计算正样本处的bce loss，注意label=score，而不是1
    loss[pos, pos_label] = F.binary_cross_entropy_with_logits(
        pred[pos, pos_label], score[pos],
        reduction='none') * scale_factor.abs().pow(beta)

    loss = loss.sum(dim=1, keepdim=False)  # (N,class_num)在1维度直接求和，变成(N,)
    return loss

def distribution_focal_loss(pred, label):
    """
    Args:
        pred (torch.Tensor): Predicted general distribution of bounding boxes
            (before softmax) with shape (N, n+1), n is the max value of the
            integral set `{0, ..., n}` in paper.
        label (torch.Tensor): Target distance label for bounding boxes with
            shape (N,).
    Returns:
        torch.Tensor: Loss tensor with shape (N,).
    """
    dis_left = label.long()  # 坐标整数值
    dis_right = dis_left + 1  # 右边整数值
    # 线性权重
    weight_left = dis_right.float() - label
    weight_right = label - dis_left.float()
    # 两个bce loss，并且加权，促使学到的分布是双峰分布，提高优化效率
    loss = F.cross_entropy(pred, dis_left, reduction='none') * weight_left \
           + F.cross_entropy(pred, dis_right, reduction='none') * weight_right
    return loss