锚框检测器损失函数异常排查：分类损失无变化

问题背景 #

我正在梳理带锚框（anchors）的检测器模型的学习过程，目前遇到了损失函数相关问题。模型的分类头、回归头以及锚框生成模块均已验证正常，但尝试多种损失函数变体后仍未解决问题。模型采用MobileNet作为特征骨干网络，连接两个适配头以输出回归和分类结果。

初始损失类代码 #

class Losses(nn.Module):
    def __init__(self, cls_loss=nn.BCEWithLogitsLoss(), reg_loss=nn.SmoothL1Loss(), threshold=0.1, alpha=1.0, beta=1.0):
        super().__init__()
        self.threshold = threshold
        self.cls_loss = cls_loss
        self.reg_loss = reg_loss
        self.alpha = alpha
        self.beta = beta

    def forward(self, bboxes, pred_offsets, predictions, groundtruth):
        """
        bboxes - bounding boxes, shape [bs, X, 4]
        pred_offsets - offset values, shape [bs, X, 4],
        predictions - prediction of target object in bbox, shape [bs, X], 
        groundtruth - real boxes, shape [bs, y, 4]. The structure is [xc, yc, w, h]
        """
        assert bboxes.shape[0] == groundtruth.shape[0], "Mismatch in batch sizes"

        distance, probs, giou_losses = [], [], []
        for i in range(bboxes.shape[0]):
            dist, prob, giou = self._overlap(bboxes[i], groundtruth[i])
            distance.append(dist)
            probs.append(prob)
            giou_losses.append(giou)

        distance = torch.stack(distance)
        probs = torch.stack(probs)
        giou_loss = torch.stack(giou_losses).mean()  

        # Classification loss
        cls_loss = self.cls_loss(predictions, probs)
        # Regression loss
        reg_loss = self.reg_loss(distance, torch.zeros_like(predictions[0]))
        # Weighted total loss
        total_loss = self.alpha * cls_loss + self.beta * reg_loss + giou_loss
        return total_loss, cls_loss, reg_loss, giou_loss

    def _overlap(self, bboxes, groundtruth):
        distance = torch.zeros(len(bboxes)).to(bboxes.device)
        probs = torch.zeros(len(bboxes)).to(bboxes.device)
        giou_loss = torch.zeros(len(bboxes)).to(bboxes.device)

        for gt in groundtruth:
            if torch.any(gt[:4] != 0):
                x1, y1, x2, y2 = to_corners(gt[:4])
                bx1, by1, bx2, by2 = to_corners_batch(bboxes)
                iou, dist = self._intersect(x1, y1, x2, y2, bx1, by1, bx2, by2)
                giou = self.generalized_iou(
                    torch.stack([bx1, by1, bx2, by2], dim=-1),
                    torch.tensor([x1, y1, x2, y2], device=bboxes.device)
                )
                mask = (iou >= self.threshold)
                probs[mask] = iou[mask]  # Use IoU as probabilities
                distance[mask] = dist[mask]
                giou_loss[mask] = (1 - giou[mask])  # GIoU values

        probs[probs == 0] = 0.01  # small noise
        return distance, probs, giou_loss

    @staticmethod
    def generalized_iou(bboxes1, bboxes2):
        x1_min = torch.min(bboxes1[:, 0], bboxes2[0])
        y1_min = torch.min(bboxes1[:, 1], bboxes2[1])
        x2_max = torch.max(bboxes1[:, 2], bboxes2[2])
        y2_max = torch.max(bboxes1[:, 3], bboxes2[3])

        inter_x1 = torch.max(bboxes1[:, 0], bboxes2[0])
        inter_y1 = torch.max(bboxes1[:, 1], bboxes2[1])
        inter_x2 = torch.min(bboxes1[:, 2], bboxes2[2])
        inter_y2 = torch.min(bboxes1[:, 3], bboxes2[3])

        inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)
        bbox1_area = (bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1])
        bbox2_area = (bboxes2[2] - bboxes2[0]) * (bboxes2[3] - bboxes2[1])

        union_area = bbox1_area + bbox2_area - inter_area
        iou = inter_area / (union_area + 1e-6)

        # Enclosing box area
        enclosing_area = (x2_max - x1_min) * (y2_max - y1_min)
        giou = iou - (enclosing_area - union_area) / (enclosing_area + 1e-6)
        return giou

    def _intersect(self, x1, y1, x2, y2, bx1, by1, bx2, by2):
        x_inter_min = torch.max(x1, bx1)
        y_inter_min = torch.max(y1, by1)
        x_inter_max = torch.min(x2, bx2)
        y_inter_max = torch.min(y2, by2)

        inter_width = torch.clamp(x_inter_max - x_inter_min, min=0)
        inter_height = torch.clamp(y_inter_max - y_inter_min, min=0)
        intersection_area = inter_width * inter_height

        union_area = ((x2 - x1) * (y2 - y1)) + ((bx2 - bx1) * (by2 - by1)) - intersection_area
        iou = intersection_area / (union_area + 1e-6)

        dist = torch.sqrt((bx2 - x2) ** 2 + (by2 - y2) ** 2) + torch.sqrt((bx1 - x1) ** 2 + (by1 - y1) ** 2)
        return iou, dist


def to_corners_batch(batch):
    cx, cy, w, h = batch[..., 0], batch[..., 1], batch[..., 2], batch[..., 3]
    bx1 = cx - w / 2
    by1 = cy - h / 2
    bx2 = cx + w / 2
    by2 = cy + h / 2
    return bx1, by1, bx2, by2


def to_corners(box):
    cx, cy, w, h = box[0], box[1], box[2], box[3]
    bx1 = cx - w / 2
    by1 = cy - h / 2
    bx2 = cx + w / 2
    by2 = cy + h / 2
    return bx1, by1, bx2, by2

修改后的损失函数代码 #

def forward(self, anchors, preds, gt_bboxes, offset):
    probs = preds['probs']  
    bboxes = preds['bboxes']  
    device = anchors.device

    cls_pos_targets = torch.zeros_like(probs, device=device)  
    cls_neg_targets = torch.zeros_like(probs, device=device)  
    reg_targets = torch.zeros_like(bboxes, device=device)   
    reg_mask = torch.zeros_like(probs, device=device)  

    iou_loss = 0
    reg_loss = 0
    cls_loss = 0
    
    bs = anchors.shape[0]
    for b in range(bs):
        gt = gt_bboxes[b]  
        gt = gt[~torch.all(gt == 0, dim=-1)]  
        if len(gt) == 0:
            continue  # if we do not have groundtruth
            
        ious = self.compute_iou(bboxes[b], gt)  # calc IoU between bboxes and groundtruth

        max_ious, gt_indices = ious.max(dim=1) 
        
        # Positive boxes
        pos_mask = max_ious >= self.pos_trh
        iou_loss += (1 - max_ious[pos_mask].sum()/pos_mask.sum()) if pos_mask.sum() > 0 else 0
        neg_mask = max_ious < self.neg_trh
        # update pos and neg targets
        cls_pos_targets[b, pos_mask] = 1
        cls_neg_targets[b, neg_mask] = 1

        # regression targets (use only positive anchors, boxes and groundtruth)
        if pos_mask.sum() > 0:
            pos_anchors = anchors[b][pos_mask]
            pos_gt = gt[gt_indices[pos_mask]]
            reg_targets[b, pos_mask] = self.box_to_deltas(pos_anchors, pos_gt) # calculate the required offset between the anchors and their groundtruth
            pred_delta = self.box_to_deltas(pos_anchors, bboxes[b, pos_mask]) # predict offset between the anchors and bboxes
            reg_loss += F.smooth_l1_loss(pred_delta, reg_targets[b, pos_mask], reduction="sum")/max(1, pos_mask.sum()) # calculate regression loss

    # classification losses. Most probably the error is here
    total = cls_pos_targets.sum() + cls_neg_targets.sum()
    pos_count = cls_pos_targets.sum().clamp(min=1)
    neg_count = cls_neg_targets.sum().clamp(min=1)
    total_count = pos_count + neg_count
    cls_pos_loss = focal(probs, cls_pos_targets).sum()
    cls_neg_loss = focal(-probs, cls_neg_targets).sum()/neg_count
    cls_loss = cls_pos_loss + cls_neg_loss
    
    # Total loss
    total_loss = cls_loss/total + self.lambda_reg * reg_loss 
    # print(f'classification loss {cls_loss/total} and regression loss {reg_loss}, iou loss = {iou_loss}')
    return total_loss

当前现象与疑问 #

分类损失几乎没有变化，而回归损失如预期下降，同时IoU损失略有下降。分类部分输出的prob logits在训练过程中没有变化，但各层梯度正常更新，推测问题出在分类损失计算环节。

内容的提问来源于stack exchange，提问作者Johan