YOLOX改进--损失函数(FocalLoss)

YOLOX官方代码改进–损失函数

---

一、置信度预测损失改进

1. FocalLoss

二元交叉熵损失BCELoss替换为FocalLoss
(1). 在YOLOX-main/yolox/models/losses.py中续写一个 FocalLoss的类
(引用的是yolov5的FocalLoss代码实现,gamma=1.5, alpha=0.25)

class FocalLoss(nn.Module):
    # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
    def __init__(self, reduction="none", gamma=1.5, alpha=0.25):
        """Initializes FocalLoss with specified loss function, gamma, and alpha values; modifies loss reduction to
        'none'.
        """
        super().__init__()
        self.loss_fcn = nn.BCEWithLogitsLoss(reduction="none")  # must be nn.BCEWithLogitsLoss()
       # self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = reduction
       # self.loss_fcn.reduction = "none"  # required to apply FL to each element

    def forward(self, pred, true):
        """Calculates the focal loss between predicted and true labels using a modified BCEWithLogitsLoss."""
        loss = self.loss_fcn(pred, true)
        # p_t = torch.exp(-loss)
        # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability

        # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
        pred_prob = torch.sigmoid(pred)  # prob from logits
        p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = (1.0 - p_t) ** self.gamma
        loss *= alpha_factor * modulating_factor

        if self.reduction == "mean":
            return loss.mean()
        elif self.reduction == "sum":
            return loss.sum()
        else:  # 'none'
            return loss

(2).在./yolox/models/yolo_head.py中追加调用FocalLoss

from .losses import IOUloss, FocalLoss

(3).在./yolox/models/yolo_head.py中实例化FocalLoss
（位置在125-131行左右）

 self.bcewithlog_loss = nn.BCEWithLogitsLoss(reduction="none")
 self.iou_loss = IOUloss(reduction="none")
 self.focal_loss=FocalLoss(reduction="none") ##FocalLoss

(4).在./yolox/models/yolo_head.py中修改loss_obj
找到置信度预测损失计算位置loss_obj，并进行替换（位置在386-395行左右）

    #   loss_obj = (
    #      self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets)
    #  ).sum() / num_fg
        loss_obj = (
            self.focal_loss(obj_preds.view(-1, 1), obj_targets)
        ).sum() / num_fg

参考

1. u版官方Yolov5的Focal_loss代码
2. Focal loss原理&代码（引用YOLOv5）
3. Focal Loss正确复现代码
4. Focal Loss论文及官方代码
5. Focal Loss官方代码