如何解决PyTorch二分类任务中数据加载器带来的类别不平衡问题

背景 #

我用PyTorch构建了二分类任务的合成数据集及对应的数据加载器，其中0类样本占比约20%，1类样本占比约80%。模型训练后只会预测占比80%的类别，原因是训练时接触1类样本的时间占比达80%。

问题1：获取数据加载器的数据后，该如何处理这种类别不平衡问题？ #

在不能修改原始数据集的前提下，可通过以下几种方式处理：

1. 加权随机采样 #

给少数类样本设置更高的采样权重，让DataLoader生成batch时更大概率选中少数类样本，平衡每个batch内的类别比例。具体实现用WeightedRandomSampler：

import torch
from torch.utils.data import DataLoader, WeightedRandomSampler

# 统计数据集内的类别分布
dataset = MyDataset()
labels = [label for _, label in dataset]
class_counts = torch.bincount(torch.tensor(labels))
# 计算采样权重：权重与类别占比成反比
class_weights = 1. / class_counts.float()
# 给每个样本分配对应类别的权重
sample_weights = class_weights[torch.tensor(labels)]
# 创建加权采样器，replacement=True允许重复采样少数类
sampler = WeightedRandomSampler(sample_weights, num_samples=len(dataset), replacement=True)
# 用采样器构建数据加载器
dl = DataLoader(dataset, batch_size=1000, sampler=sampler)

2. 损失函数加权 #

给少数类样本的损失设置更高权重，让模型对少数类的错误预测付出更大代价，这个方法可与加权采样结合使用，效果更稳定，具体实现见后续问题说明。

3. 优化评估指标 #

不要仅依赖准确率（准确率会被多数类主导），改用F1-score、AUC-ROC、召回率等对不平衡数据更友好的指标，这些指标能真实反映模型对少数类的识别能力。

问题2：BCELoss能否应对这种类别不平衡场景？ #

默认的BCELoss不行——它平等对待所有样本的损失，多数类样本数量多，总损失占比更高，模型会自然偏向多数类。但通过设置weight参数，BCELoss可以有效适配类别不平衡场景，本质是给不同类别的样本损失加权，让少数类的损失更受重视。

补充问题：BCELoss的weight参数能否用于处理类别不平衡问题？（无法修改原始数据集） #

完全可以，这是处理类别不平衡的常用方法之一。

具体实现步骤： #

1. 计算类别权重：权重与类别占比成反比，比如0类占20%、1类占80%，则0类权重为1/0.2=5，1类权重为1/0.8=1.25，这样两类样本的总损失期望会趋于均衡。
2. 训练时，给每个batch内的样本匹配对应类别的权重，传入BCELoss的weight参数。

代码示例（推荐用BCEWithLogitsLoss，它整合了sigmoid和BCELoss，数值更稳定）：

import torch
from tqdm import tqdm

# 计算类别权重
class_counts = torch.tensor([0.2*50000, 0.8*50000])  # 对应0类、1类的样本数
class_weights = 1. / class_counts.float()
class_weights = class_weights.to(DEVICE)

# 使用BCEWithLogitsLoss替代单独的sigmoid+BCELoss
loss_fn = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-6)

losses = []
for epoch in range(num_epochs):
    print(f"EPOCH: {epoch}")
    for data, label in tqdm(dl, total=len(dl)):
        data = data.to(DEVICE).float()
        label = label.to(DEVICE).float()
        # 生成当前batch的样本权重
        batch_weights = class_weights[label.long()]
        optimizer.zero_grad()
        preds = model(data)
        # 传入权重参数计算损失
        loss = loss_fn(preds, label, weight=batch_weights)
        loss.backward()
        losses.append(loss.detach().cpu().numpy())
        optimizer.step()

注意事项： #

• 权重可根据实际需求调整，若少数类误判代价更高，可进一步放大其权重。
• 加权采样+损失加权的组合策略，通常比单一方法效果更好。

内容的提问来源于stack exchange，提问作者Saeed