@LOSSES.register_module() class LaneCustomWeightedLoss(nn.Module):def __init__(self,negative_ratio=3,loss_weight=1.0,type_num=3, # Assuming there are 11 classesspecial_classes=[4],special_weight=5)super(LaneCustomWeightedLoss, self).__init__()self.type_num = type_numself.negative_ratio = negative_ratioself.loss_weight = loss_weightself.special_classes = special_classesself.special_weight = special_weightdef forward(self, pred_seg, gt_seg_cur, quality_weight): # pred_seg：（bs, num_cls, H, W）B, H, W = gt_seg_cur.shapegt_seg_cur = gt_seg_cur.reshape(-1)pred_seg = pred_seg.permute(0, 2, 3, 1).reshape(-1, self.type_num)pmask = (gt_seg_cur > 0) # 创建正样本掩码，标记真实标签中大于0的位置（即有车道线的位置）fpmask = pmask.float()n_pos = torch.sum(fpmask)nmask = (gt_seg_cur == 0) # 创建负样本掩码，标记真实标签中等于0的位置（即无车道线的背景位置）fnmask = nmask.float()max_neg_entries = torch.sum(fnmask)n_neg = n_pos * self.negative_ratio # 计算需要选择的负样本数量，按照正负样本比例（默认为1:3）n_neg = n_neg if n_neg < max_neg_entries else max_neg_entriesn_neg = n_neg if n_neg > 0 else torch.ones_like(n_neg)n_neg_int = n_neg.int()bg_predict = pred_seg[:, 0] # 获取预测结果中背景类别的预测值val, indxes = torch.topk(bg_predict[nmask], k=n_neg_int, largest=False) # 在负样本位置中找到预测值最小的 n_neg_int 个样本（最难的负样本）max_hard_pred = val[-1]nmask_ohem = ((bg_predict <= max_hard_pred) * nmask) # 通过OHEM策略创建最终的负样本掩码，选择最难的负样本fnmask_ohem = nmask_ohem.float()log_p = F.log_softmax(pred_seg, dim=-1)q_log_p = -log_p.gather(dim=1, index=gt_seg_cur.unsqueeze(1).long()).squeeze(1) # 提取真实类别对应的对数概率，并取负号得到交叉熵损失# Apply special weight to specific classesspecial_mask = torch.zeros_like(gt_seg_cur, dtype=torch.float)for cls in self.special_classes:special_mask += (gt_seg_cur == cls).float()weighted_q_log_p = q_log_p * (1 + self.special_weight * special_mask) # 为特殊类别应用额外权重，普通类别权重为1，特殊类别权重为 1+special_weight# 根据质检结果改变loss权重assert len(quality_weight) == Bquality_weight = torch.tensor(quality_weight, device=pred_seg.device)[:, None, None].expand(B, H, W)quality_weight = quality_weight.reshape(-1)weighted_q_log_p *= quality_weighttotal_cross_pos = torch.sum(weighted_q_log_p * fpmask) # 计算加权后的正样本交叉熵损失总和total_cross_neg = torch.sum(weighted_q_log_p * fnmask_ohem) # 计算加权后的负样本交叉熵损失总和（仅包含OHEM选中的困难负样本）total_cross_pos = total_cross_pos / n_pos if n_pos > 0 else torch.zeros_like(total_cross_pos)total_cross_neg = total_cross_neg / n_negtotal_cross_pos = total_cross_pos * self.loss_weighttotal_cross_neg = total_cross_neg * self.loss_weightreturn total_cross_pos + total_cross_neg, pmask.view(B,H,W)
