Ohem Cross Entropy Loss实战：如何用PyTorch解决图像分割中的样本不平衡问题

在图像分割的实际项目中，我们常常会遇到一个令人头疼的问题：模型总是对那些数量庞大的“简单”背景区域学得很好，而对真正重要的、但数量稀少的“困难”目标区域（比如医学图像中的病灶、自动驾驶场景中的行人）表现不佳。这背后，就是典型的样本不平衡在作祟。传统的交叉熵损失函数对所有像素一视同仁，导致模型优化方向被多数类“绑架”，最终在关键指标上表现平平。

今天，我们不谈空洞的理论，直接从代码和实战出发，聊聊如何用PyTorch实现Ohem Cross Entropy Loss，这个专治样本不平衡的“靶向药”。它的核心思想非常直观：在训练过程中，动态地、有选择地“教训”那些模型还搞不定的“困难样本”，而对于已经学得很好的“简单样本”，则适当“放水”。这种方法能让你的模型把有限的注意力，精准地投入到最需要提升的地方。

无论你是在做医学影像分析、自动驾驶感知，还是卫星图像解译，只要你的分割任务中存在前景与背景像素数量严重不均，或者不同类别目标大小差异巨大的情况，这篇文章提供的思路和代码，或许就能成为你突破模型性能瓶颈的那把钥匙。

1. 理解OHEM：不止是“困难样本挖掘”

在深入代码之前，我们有必要先跳出“在线困难样本挖掘”这个字面定义，从更本质的优化视角来理解OHEM。它解决的，其实是损失函数梯度贡献不均衡的问题。

想象一下，在一张城市街景的分割图中，天空和道路的像素可能占了80%以上，而行人和车辆只占不到5%。使用标准交叉熵损失时，模型即使把所有行人和车辆都预测错了，但只要把天空和道路预测对，整体损失依然可以很低。模型自然会选择“躺平”，专注于学习多数类。OHEM通过一个巧妙的机制，打破了这种“沉默的大多数”对训练过程的统治。

OHEM的核心运作机制可以概括为两个步骤：

  1. 筛选 (Selection)：在前向传播计算得到每个像素（或样本）的初始损失后，根据一个预设的阈值（例如，对应预测概率为0.7），将所有损失高于此阈值的样本标记为“困难样本”。这些样本通常是模型当前预测置信度低、容易出错的。
  2. 重加权 (Reweighting)：在反向传播时，只使用这批筛选出的困难样本的损失来计算梯度，或者给予它们更高的权重。简单样本的梯度被直接置零或大幅衰减。

注意：这里的“困难”是动态变化的。随着模型训练，昨天还很难分的样本，今天可能就变简单了。OHEM能自动适应这种变化，始终聚焦于当前的“难点”。

为了更清晰地对比OHEM与标准交叉熵的差异，我们可以看下面这个表格：

特性维度标准交叉熵损失 (CrossEntropyLoss)OHEM交叉熵损失 (OhemCELoss)
梯度来源所有样本（无论难易）仅困难样本（或困难样本权重极高）
关注点整体平均精度困难样本的分类性能
对不平衡数据的鲁棒性弱，易被多数类主导强，主动聚焦少数困难类
训练稳定性高，梯度来源稳定相对较低，需谨慎设置阈值和最小样本数
计算开销标准前向/反向传播需额外进行损失排序或筛选操作
典型应用场景类别均衡或对各类别错误同等关心的任务前景-背景严重不平衡、小目标检测、关键区域分割

从优化角度看，OHEM可以被视为一种自适应的、基于难度的样本重采样策略。它不是在数据层面进行过采样或欠采样，而是在损失层面进行动态调整，实现更高效的梯度更新。

2. 从零构建PyTorch版OhemCELoss

理解了原理，我们动手实现一个功能完整、适用于图像分割的OhemCELoss。我们将采用一种兼顾效率与效果的实现方式，并详细解释每一个参数和设计选择。

首先，导入必要的模块：

import torch
import torch.nn as nn
import torch.nn.functional as F

接下来是损失函数类的定义。我们的实现将包含几个关键设计：

class OhemCrossEntropyLoss(nn.Module):
    """
    在线困难样本挖掘交叉熵损失函数 (PyTorch实现)
    适用于图像分割任务，有效缓解类别不平衡问题。

    Args:
        thresh (float): 困难样本阈值。损失值大于 `-log(thresh)` 的样本被认定为困难样本。
        min_kept (int, optional): 每批训练中至少保留参与损失计算的像素数。默认为批大小*高*宽 // 16。
        ignore_index (int, optional): 在标签中需要忽略的索引值（如255）。默认为255。
        reduction (str, optional): 对最终损失的处理方式，'mean' 或 'none'。默认为'mean'。
    """
    def __init__(self, thresh=0.7, min_kept=None, ignore_index=255, reduction='mean'):
        super(OhemCrossEntropyLoss, self).__init__()
        # 将概率阈值转换为损失阈值。因为交叉熵损失 = -log(p)，所以p>thresh 等价于 loss < -log(thresh)
        self.thresh = -torch.log(torch.tensor(thresh, dtype=torch.float))
        self.min_kept = min_kept
        self.ignore_index = ignore_index
        self.reduction = reduction
        # 基础交叉熵损失函数，用于计算每个像素的原始损失
        self.criterion = nn.CrossEntropyLoss(ignore_index=ignore_index, reduction='none')

    def forward(self, logits, labels):
        """
        Args:
            logits (torch.Tensor): 模型预测输出，形状为 [N, C, H, W]
            labels (torch.Tensor): 真实标签，形状为 [N, H, W]，像素值为类别索引。

        Returns:
            torch.Tensor: 计算得到的OHEM损失值。
        """
        # 1. 计算每个像素的原始交叉熵损失
        # logits: (N, C, H, W) -> 需要view为 (N*H*W, C) 以适应F.cross_entropy，但我们使用已初始化的criterion
        # labels: (N, H, W)
        pixel_losses = self.criterion(logits, labels)  # 形状: (N, H, W) 或展平后为 (N*H*W,)

        # 2. 展平损失和标签，便于后续操作
        n, h, w = labels.shape
        pixel_losses = pixel_losses.view(-1)  # 形状: (N*H*W,)
        valid_mask = (labels.view(-1) != self.ignore_index)  # 忽略无效像素（如边界或padding）

        # 3. 获取有效像素的损失
        valid_losses = pixel_losses[valid_mask]
        if valid_losses.numel() == 0:
            # 如果没有有效像素，返回零损失
            return torch.tensor(0.0, device=logits.device, requires_grad=True)

        # 4. 确定本轮训练需要保留的最小像素数
        if self.min_kept is None:
            # 默认策略：至少保留有效像素数的 1/16，确保训练持续进行
            min_kept = max(1, valid_losses.numel() // 16)
        else:
            min_kept = self.min_kept
        # 确保最小保留数不超过有效像素总数
        min_kept = min(min_kept, valid_losses.numel())

        # 5. 困难样本挖掘：选择损失最大的前K个像素
        # 这里使用 `topk` 操作，直接获取损失值最大的前 min_kept 个样本。
        # 同时，我们也应用阈值：只有损失大于 self.thresh 的样本才被认为是“困难”的。
        # 但为了确保有足够的样本进行训练，我们优先保证数量 (min_kept)。
        loss_thresh, _ = torch.topk(valid_losses, k=min_kept, sorted=True)
        # 最终阈值取“第min_kept大的损失值”和“预设概率阈值对应的损失值”中的较大者。
        # 这保证了至少有min_kept个样本被选中，且选中的样本损失都至少达到了某个“困难”水平。
        loss_thresh = max(loss_thresh[-1], self.thresh.to(valid_losses.device))

        # 6. 生成最终用于计算损失的掩码
        # 有效像素中，损失大于等于最终阈值的，才是本轮参与梯度计算的困难样本。
        hard_sample_mask = (pixel_losses >= loss_thresh) & valid_mask
        # 如果没有任何像素满足条件（理论上不会发生，因为min_kept>=1），则回退到所有有效像素
        if torch.nonzero(hard_sample_mask).numel() == 0:
            hard_sample_mask = valid_mask

        # 7. 计算最终损失
        # 只对困难样本计算损失
        final_loss = pixel_losses[hard_sample_mask]
        if self.reduction == 'mean':
            return final_loss.mean()
        elif self.reduction == 'sum':
            return final_loss.sum()
        else: # 'none'
            # 返回每个困难样本的损失，形状不规则，通常不推荐
            return final_loss

这个实现有几个关键点值得深入讨论：

  • 阈值 (thresh) 的动态性：代码中最终的 loss_thresh 是动态确定的。它取 预设阈值 和 第min_kept大的损失值 之间的较大者。这意味着：
    • 当模型很差时，很多样本的损失都很大，第min_kept大的损失值可能远高于预设阈值，此时阈值较高，筛选非常严格。
    • 当模型很好时，困难样本变少，第min_kept大的损失值可能低于预设阈值。此时，为了满足min_kept的要求，阈值会自动降低，确保有足够样本驱动训练，防止训练停滞。
  • 最小保留数 (min_kept) 的安全网作用：这是OHEM稳定训练的关键。如果没有这个设置，在训练后期，可能几乎没有像素的损失能超过固定阈值，导致有效梯度为零，训练提前收敛到一个次优点。min_kept 保证了训练信号的持续性。
  • 忽略标签 (ignore_index) 的处理：在分割数据集中，标签为特定值（常为255）的像素通常表示边界或无效区域。我们的实现正确地在筛选困难样本前将其排除，避免了噪声干扰。

你可以像使用任何标准PyTorch损失函数一样使用它：

# 初始化
criterion = OhemCrossEntropyLoss(thresh=0.7, ignore_index=255)

# 在训练循环中
for images, masks in dataloader:
    images, masks = images.to(device), masks.to(device)
    outputs = model(images)
    loss = criterion(outputs, masks)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

3. 参数调优与实战技巧：让OHEM发挥最大效力

OHEM虽然强大，但它引入了新的超参数。用不好，效果可能还不如标准损失函数。下面结合我的项目经验，分享一些调参心得和实战技巧。

3.1 核心参数详解与调优指南

  • thresh (阈值): 这是最重要的参数。它决定了“困难”的门槛。

    • 含义：对应的是模型预测正确类别的概率。thresh=0.7意味着，对于某个像素，如果模型预测其真实类别的概率高于70%，就认为是简单样本，可能被忽略。
    • 设置策略：
      • 初始值：通常从 0.7 开始尝试。这是一个经验性的平衡点。
      • 调高 (>0.7)：会使筛选更严格，只有模型非常不确定的样本才参与训练。这能让模型更专注于最难的案例，但可能导致训练不稳定、收敛慢，因为梯度信号变弱、变稀疏。适用于数据噪声小、你特别关心最难样本性能的场景。
      • 调低 (<0.7)：会让更多样本被认定为“困难”，训练更接近标准交叉熵，但依然对简单样本有所抑制。这能稳定训练，加快初期收敛。适用于数据噪声较大、或者你希望OHEM效应温和一些的情况。
    • 监控方法：训练时，可以打印或记录每一批中被选为困难样本的像素比例。这个比例在训练初期应该较高（如30%-50%），随着模型变好，逐渐下降并稳定在一个较低值（如5%-15%）。如果比例始终很高或很低，可能需要调整thresh。

  • min_kept (最小保留数): 训练稳定性的“保险丝”。

    • 含义：无论如何，每批至少要保证有这么多像素参与损失计算。
    • 设置策略：
      • 默认值：有效像素总数 // 16 是一个很好的起点。它大约对应总像素的6%。
      • 调高：会增加每次迭代的梯度估计的稳定性，但会稀释OHEM的效果（因为混入了更多相对简单的样本）。
      • 调低：会让OHEM效应更强，但训练可能因梯度方差过大而震荡。一般不建议设得过低，尤其是在批大小（batch size）本身就不大的情况下。
    • 与批大小的关系：min_kept是一个绝对值。当你的批大小增大时，有效像素总数增多，默认的 //16 也会相应增大，这是合理的。如果你固定了min_kept，当换用更大的批大小时，可能需要适当调高它。

3.2 与其他技术联用的策略

OHEM很少单独使用，结合其他技术能产生“1+1>2”的效果。

  • 与类别权重 (Class Weight) 结合：这是最常用的组合。类别权重在数据层面给少数类更高的惩罚，而OHEM在样本难度层面给困难样本更高的关注。两者从不同维度攻击不平衡问题。

    # 假设我们有一个3类的分割任务，背景(0)、类别1(1)、类别2(2)
# 计算各类别像素频率的倒数作为权重（一种常见方法）
# weight_for_class = 1.0 / (class_frequency + epsilon)
class_weights = torch.tensor([0.1, 1.0, 3.0], device=device) # 示例权重，类别2最稀有，权重最高

# 修改OHEM初始化，将权重传入基础交叉熵
self.criterion = nn.CrossEntropyLoss(weight=class_weights, ignore_index=ignore_index, reduction='none')

    提示：可以先只用类别权重训练几个epoch，让模型有一个初步的类别感知，再加入OHEM，这样训练会更平滑。

  • 与Focal Loss的对比与选择：

    • Focal Loss：通过 (1-p)^gamma 的调制因子，连续地降低简单样本的损失贡献。它是一个“软”筛选。
    • OHEM：通过阈值，硬性地忽略简单样本的梯度。它是一个“硬”筛选。
    • 如何选：在许多实验中，OHEM的表现略优于或与Focal Loss相当。OHEM的计算更简单直观，调参逻辑（阈值）也更直接。Focal Loss多了一个 gamma 参数需要调节。如果你追求极致的性能，可以尝试将两者结合（即使用Focal Loss作为OHEM内部的基损失），但这会进一步增加复杂性。

  • 学习率调整：由于OHEM使得每次参与计算的样本是高度变化的、困难的，这相当于给优化过程引入了更大的噪声。可以考虑使用更小的初始学习率，或者使用带有热身 (Warmup) 策略的学习率调度器，让模型在初期更平稳地适应这种噪声。

4. 实战案例：在医学图像分割中应用OHEM

让我们看一个具体的例子：在视网膜血管分割数据集（如DRIVE）上应用OHEM。这个任务的挑战是，血管像素（前景）只占整张图像的不到10%，是典型的极度不平衡问题。

项目背景与基线： 我们使用一个经典的U-Net作为基线模型，使用标准交叉熵损失训练。评估指标除了整体准确率(Accuracy)，我们更关注敏感度(Sensitivity) 和 交并比(IoU)，因为它们更能反映模型对稀少血管的捕捉能力。

基线模型结果可能如下：

  • 准确率(Accuracy): ~0.97 （很高，因为背景预测对了就行）
  • 血管IoU: ~0.55
  • 敏感度(Sensitivity): ~0.65 （很多细血管被漏检）

引入OHEM的改进步骤：

  1. 损失函数替换：将 nn.CrossEntropyLoss 替换为我们自定义的 OhemCrossEntropyLoss。初始参数设为 thresh=0.7, min_kept=None（使用默认值）。

  2. 训练过程观察：

    • 训练初期，损失值会比使用标准损失时更高且波动更大，这是正常的，因为模型正在努力拟合那些它之前忽略的困难血管像素。
    • 通过添加简单的日志，观察困难像素比例：
      # 在OhemCELoss的forward函数末尾添加（调试用）
hard_ratio = hard_sample_mask.float().mean().item()
# 可以定期打印或使用TensorBoard记录
      你会发现这个比例从开始的~40%逐渐下降到~10%并趋于稳定。

  3. 参数微调：

    • 训练一段时间后，如果验证集IoU提升缓慢，可以尝试将 thresh 从0.7略微提高到0.75或0.8，迫使模型去挖掘那些“更难”的血管末梢和低对比度区域。
    • 如果训练出现不稳定（损失剧烈震荡），可以尝试略微调大 min_kept（例如从 //16 改为 //12），或者减小学习率。

  4. 结合数据增强：对于小目标（如细血管），随机裁剪、缩放、旋转等几何增强非常重要。OHEM帮助模型关注困难样本，而数据增强则创造了更多样的困难样本变体，两者相辅相成。

改进后的典型结果：

  • 准确率(Accuracy): 可能略微下降至 ~0.96 （因为模型不再一味讨好背景）
  • 血管IoU: 提升至 ~0.65 （显著提升）
  • 敏感度(Sensitivity): 提升至 ~0.78 （漏检大幅减少）

可视化分割结果对比会非常明显：使用OHEM后，预测的血管网络更连续，尤其是那些纤细的末梢血管被更好地检测出来。

在另一个城市街景分割项目中，OHEM同样显著提升了对远处小车辆、行人的分割精度。它的价值在于，它不仅仅是一个损失函数，更是一种引导模型优化方向的策略。它告诉优化器：“不要满足于在简单样本上刷分，要去攻克那些真正的难关。”

最后，记住一点：OHEM不是银弹。在类别相对均衡的数据集上，它可能带来不必要的复杂度。但在那些“关键区域只占少数”的视觉任务里，当你觉得模型性能遇到瓶颈，常规方法无力时，亲手实现并调试一个OHEM损失，很可能会给你带来意想不到的突破。调试的关键在于耐心观察训练动态，理解 thresh 和 min_kept 如何影响模型的学习焦点，然后像打磨一件工具一样，将它调整到最适合你当前任务的状态。

Logo
DAMO开发者矩阵

DAMO开发者矩阵，由阿里巴巴达摩院和中国互联网协会联合发起，致力于探讨最前沿的技术趋势与应用成果，搭建高质量的交流与分享平台，推动技术创新与产业应用链接，围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐

cover

【厦门大学主办 | 往届均已EI稳定检索！】第六届人工智能、机器人和通信国际会议（ICAIRC 2026）

avatar DAMO开发者矩阵

企业微信API接口如何打造外部群智能助手？从群消息处理到自动响应

外部群智能助手这套东西，难点不在 AI 多能聊，在群消息识别、角色判断、上下文管理、响应策略、防刷屏这些工程细节。每一项都不深奥，但少做一项机器人在群里就是个灾难。这套搭扎实，机器人在客户群里该说话时说话、不该说时闭嘴、说一次说清楚——才真是个群助手而不是群噪音。

avatar DAMO开发者矩阵

2026微信机器人怎么搭建？手把手从0跑通收发消息（含完整代码）

很多人想做微信机器人，一上来就被各种框架、协议、SDK绕晕。其实跑通一个能用的机器人，最短路径只有三件事：让号在线、能发消息、能收消息。这篇用Python手把手走一遍，复制改Token就能跑。接口路径、参数与回调字段以api文档weiti.apifox.cn 为准。

avatar DAMO开发者矩阵