mmrotate代码解读之RotatedAnchorHead—loss_single()函数

loss_single()函数功能

它用于计算单一尺度级别特征图上的损失

参数解析：

1.cls_score (torch.Tensor): 单个尺度的类别分数，形状为 (N, num_anchors * num_classes, H, W)，其中N是批次大小，num_anchors是每个位置的锚框数量，num_classes是类别数量，H和W是特征图的高和宽。

2.bbox_pred (torch.Tensor): 单个尺度的边界框预测，形状为 (N, num_anchors * 5, H, W)。5代表。x,y,w,h,角度

3.anchors (torch.Tensor): 每个尺度级别的锚框参考，形状为 (N, num_total_anchors, 5)。

4.labels (torch.Tensor): 每个锚框的标签，形状为 (N, num_total_anchors)。

5.label_weights (torch.Tensor): 每个锚框的标签权重，形状为 (N, num_total_anchors)。

6.bbox_targets (torch.Tensor): 每个锚框的边界框回归目标，形状为 (N, num_total_anchors, 5)。它是由标签分配得来的

7.bbox_weights (torch.Tensor): 每个锚框的边界框回归损失权重，形状为 (N, num_total_anchors, 5)。

8.num_total_samples (int): 如果进行采样，则为总锚框数；否则，它是正锚框的数量。

函数流程：

一.分类损失（loss_cls）的计算:

1.labels和label_weights被重塑为一维张量，cls_score通过permute和reshape操作被重塑，以匹配labels的维度。
2.使用self.loss_cls方法计算分类损失，传入重塑后的cls_score、labels、label_weights，以及avg_factor作为平均因子（通常是总样本数或正样本数）。

二.回归损失（loss_bbox）的计算:

1.bbox_targets和bbox_weights被重塑为与bbox_pred相匹配的形状。
2.如果设置了self.reg_decoded_bbox为True，则使用self.bbox_coder.decode方法对bbox_pred进行解码，这通常是将模型预测转换为实际的边界框坐标。
3.使用self.loss_bbox方法计算回归损失，传入bbox_pred、bbox_targets、bbox_weights，以及avg_factor。

代码详细解读：

# 计算单一尺度级别特征图损失
def loss_single(self, cls_score, bbox_pred, anchors, labels, label_weights,
                bbox_targets, bbox_weights, num_total_samples):

    # 将标签和标签权重重塑为一维张量，以便进行计算
    labels = labels.reshape(-1)
    label_weights = label_weights.reshape(-1)
    
    # 调整分类得分的形状以匹配标签，从(N, num_anchors * num_classes, H, W)变为(-1, self.cls_out_channels)
    cls_score = cls_score.permute(0, 2, 3, 1).reshape(-1, self.cls_out_channels)
    
    # 调用类别损失函数计算分类损失
    loss_cls = self.loss_cls(
        cls_score, labels, label_weights, avg_factor=num_total_samples)
    
    # 将边界框目标和权重重塑为匹配的形状以进行回归损失计算
    bbox_targets = bbox_targets.reshape(-1, 5)
    bbox_weights = bbox_weights.reshape(-1, 5)
    
    # 调整边界框预测的形状以进行计算
    bbox_pred = bbox_pred.permute(0, 2, 3, 1).reshape(-1, 5)
    
    # 如果设置了解码边界框，对预测的边界框进行解码
    if self.reg_decoded_bbox:
        anchors = anchors.reshape(-1, 5)
        bbox_pred = self.bbox_coder.decode(anchors, bbox_pred)

    # 调用边界框损失函数计算回归损失
    loss_bbox = self.loss_bbox(
        bbox_pred,
        bbox_targets,
        bbox_weights,
        avg_factor=num_total_samples)
    
    # 返回分类损失和回归损失
    return loss_cls, loss_bbox