1. Ohem Cross Entropy Loss 的定义

OhemCrossEntropyLoss 是一种用于深度学习中目标检测任务的损失函数，它是针对不平衡数据分布和困难样本训练的一种改进版本的交叉熵损失函数。Ohem 表示 “Online Hard Example Mining”，意为在线困难样本挖掘。在目标检测任务中，由于背景类样本通常远远多于目标类样本，导致了数据分布的不平衡问题，而且一些困难的样本对于网络的训练很有挑战性。OhemCrossEntropyLoss 就是为了解决这些问题而设计的。

这个损失函数的核心思想是在训练过程中只选择那些具有较高损失值的困难样本进行梯度更新，从而更加关注于难以分类的样本，有助于网络更好地适应这些样本，提高模型的性能。

数学上，OhemCrossEntropyLoss 的定义可以用以下公式表示：

$$\text{OhemCrossEntropyLoss}=-\frac{1}{N}\sum _{i=1}^N\{\begin{array}{ll}\text{log}(p_{\text{target}})& \text{if }{y}_{\text{target}}=1\text{ (目标类样本)}\\ \text{log}(1-p_{\text{target}})& \text{if }{y}_{\text{target}}=0\text{ (背景类样本且损失高于阈值)}\\ 0& \text{otherwise}\end{array}$$

其中，$N$ 是 Batch 中样本的数量，$p_{\text{target}}$ 是模型预测目标类的概率，$y_{\text{target}}$ 是真实标签（1 表示目标类，0 表示背景类），损失计算根据标签的情况进行不同的处理。背景类样本中损失值高于一个预定义的阈值的样本会被选中进行梯度更新，这样网络更关注于难以分类的样本，有助于提高性能。

需要注意的是，OhemCrossEntropyLoss 需要在训练过程中动态地筛选困难样本，所以相比于传统的交叉熵损失，它的计算相对复杂。但在处理不平衡数据和困难样本时，它能够提升模型的鲁棒性和泛化能力。

2. OHEM 步骤流程

1. 给 OhemCE Loss 取一个阈值 thresh：

   • 那么该像素点的预测概率 > 0.7，则该像素点可以看成是简单样本，不参与损失计算
   • 那么该像素点的预测概率 < 0.7，则该像素点可以看成是困难样本，参与损失计算

2. 确定忽略的像素点值 lb_ignore：一般我们将背景的值设置为 255，即如果像素点值的大小是 255，那么就不参与损失计算。

3. 设置最少计算的像素点个数 n_min：至少有 n_num 个像素点参与损失计算（不然网络有可能停止更新了）。

简单来说：OHEM CrossEntropy Loss 的目的是：挖掘困难样本；忽略简单样本。

3. 代码实现

import random
import numpy as np
import torch
from torch.autograd import Variable
import torch.nn as nn


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    
    
class OhemCELoss(nn.Module):
    def __init__(self, thresh, lb_ignore=255, ignore_simple_sample_factor=16):
        """
            Args:
                thresh: 阈值，超过该值则被算法简单样本 -> 不参与Loss计算
                lb_ignore: 忽略的像素值(一般255代表背景), 不参与损失的计算
                ignore_simple_sample_factor: 忽略简单样本的系数
                                                    该系数越大，最少计算的像素点个数越少
                                                    该系数越小，最少计算的像素点个数越多
        """
        super(OhemCELoss, self).__init__()
        
        """
            这里的 thresh 和 self.thresh 不是一回儿事儿
                ①预测概率 > thresh -> 简单样本
                ①预测概率 < thresh -> 困难样本
                ②损失值 > self.thresh -> 困难样本
                ②损失值 < self.thresh -> 简单
                
                ①和②其实是一回儿事儿，但 thresh 和 self.thresh 不是一回儿事儿
        """
        self.thresh = -torch.log(input=torch.tensor(thresh, requires_grad=False, dtype=torch.float))
        self.lb_ignore = lb_ignore
        self.criteria = nn.CrossEntropyLoss(ignore_index=lb_ignore, reduction='none')
        self.ignore_simple_sample_factor = ignore_simple_sample_factor
        
        """
            reduction 参数用于控制损失的计算方式和输出形式。它有三种可选的取值：
                1. 'none'：当设置为 'none' 时，损失将会逐个样本计算，返回一个与输入张量相同形状的损失张量。
                           这意味着输出的损失张量的形状与输入的标签张量相同，每个位置对应一个样本的损失值。
                2. 'mean'：当设置为 'mean' 时，损失会对逐个样本计算的损失进行求均值，得到一个标量值。
                           即计算所有样本的损失值的平均值。
                3. 'sum' : 当设置为 'sum'  时，损失会对逐个样本计算的损失进行求和，得到一个标量值。
                           即计算所有样本的损失值的总和。

            在语义分割任务中，通常使用 ignore_index 参数来忽略某些特定标签，例如背景类别。
            当计算损失时，将会忽略这些特定标签的损失计算，以避免这些标签对损失的影响。
            如果设置了 ignore_index 参数，'none' 的 reduction 参数会很有用，因为它可以让你获取每个样本的损失，包括被忽略的样本。

            总之，reduction 参数允许在计算损失时控制输出形式，以满足不同的需求。
        """
    
    def forward(self, logits, labels):
        # 1. 计算 n_min(至少算多少个像素点)
        n_min = labels[labels != self.lb_ignore].numel() // self.ignore_simple_sample_factor
        
        # 2. 使用 CrossEntropy 计算损失, 之后再将其展平
        loss = self.criteria(logits, labels).view(-1)
        
        # 3. 选出所有loss中大于self.thresh的像素点 -> 困难样本
        loss_hard = loss[loss > self.thresh]
        
        # 4. 如果总数小于 n_min, 那么肯定要保证有 n_min 个像素点的 loss
        if loss_hard.numel() < n_min:
            loss_hard, _ = loss.topk(n_min)
            
        # 5. 如果参与的像素点的个数 > n_min 个，那么这些点都参与计算
        loss_hard_mean = torch.mean(loss_hard)
        
        # 6. 返回损失的均值
        return loss_hard_mean
    
    
if __name__ == "__main__":
    setup_seed(20)
    
    # 1. 生成预测值(假设我们有两个样本，每个样本有 3 个类别，高度和宽度均为 4)
    logits = Variable(torch.randn(2, 3, 4, 4))  # [N, C, H, W], s.t. C <-> num_classes
    
    # 2. 生成真实标签(每个样本的标签是一个 4x4 的图像)
    labels = Variable(torch.randint(low=0, high=3, size=(2, 4, 4)))  # [N, H, W]
    
    # 3. 初始化：创建 OhemCELoss 的实例，阈值设置为 0.7
    ohem_criterion = OhemCELoss(thresh=0.7, lb_ignore=255, ignore_simple_sample_factor=16)
    
    # 4. 计算 Ohem 损失
    loss = ohem_criterion(logits, labels)
    
    print(f"Ohem Loss: {loss.item()}")  # Ohem Loss: 1.3310734033584595

知识来源

1. https://www.bilibili.com/video/BV12841117yo
2. https://www.bilibili.com/video/BV1Um4y1L753