别再为PyTorch数据不平衡发愁了！手把手教你用WeightedRandomSampler搞定猫狗分类

用WeightedRandomSampler解决猫狗分类中的数据不平衡问题

当你第一次尝试用PyTorch构建猫狗分类器时，可能会遇到一个令人头疼的问题：你的数据集中猫的图片有1000张，而狗的图片只有200张。这种数据不平衡会导致模型总是倾向于预测"猫"，因为这样就能获得80%的准确率——但这显然不是我们想要的结果。

1. 理解数据不平衡的影响

在机器学习中，数据不平衡是指不同类别的样本数量存在显著差异。以我们的猫狗分类为例：

• 猫的图片：1000张
• 狗的图片：200张

这种5:1的比例会导致几个问题：

1. 模型偏见：模型会倾向于预测多数类（猫），因为这样就能获得较高的准确率
2. 评估失真：即使准确率达到80%，对少数类（狗）的分类效果可能非常差
3. 训练不稳定：少数类的样本难以对模型参数产生足够影响

常见解决方案对比：

2. WeightedRandomSampler原理解析

WeightedRandomSampler是PyTorch提供的一种采样器，它通过为每个样本分配权重来解决数据不平衡问题。其核心思想是：

• 为少数类样本分配更高的权重
• 在训练过程中，根据权重随机选择样本
• 使得每个batch中的类别比例更加均衡

关键参数说明：

torch.utils.data.WeightedRandomSampler( weights, # 每个样本的权重序列 num_samples, # 要采样的总数 replacement=True # 是否允许重复采样 )

注意：weights参数对应的是每个样本的权重，而不是类别的权重。这意味着如果你的数据集有1200个样本（1000猫+200狗），weights应该是一个长度为1200的序列。

3. 实战：为猫狗数据集实现加权采样

让我们一步步实现一个完整的解决方案。

3.1 准备数据集

假设我们有一个包含猫狗图片的文件夹结构如下：

data/ train/ cat/ cat001.jpg cat002.jpg ... dog/ dog001.jpg dog002.jpg ...

首先，我们需要统计每个类别的样本数量：

from torchvision.datasets import ImageFolder import os dataset = ImageFolder('data/train') cat_count = sum([1 for _, label in dataset if label == 0]) # 假设猫是类别0 dog_count = len(dataset) - cat_count print(f"猫的图片数量: {cat_count}, 狗的图片数量: {dog_count}")

3.2 计算样本权重

我们需要为每个样本分配权重，使得少数类样本有更高的被采样概率：

import torch # 计算每个类别的权重 class_weights = [1./cat_count, 1./dog_count] # 为每个样本分配对应的类别权重 sample_weights = [0] * len(dataset) for idx, (_, label) in enumerate(dataset): sample_weights[idx] = class_weights[label] # 转换为Tensor weights = torch.DoubleTensor(sample_weights)

3.3 创建采样器和DataLoader

现在我们可以创建采样器并将其整合到DataLoader中：

from torch.utils.data import DataLoader, WeightedRandomSampler # 创建采样器 sampler = WeightedRandomSampler( weights=weights, num_samples=len(weights), # 通常与数据集大小相同 replacement=True # 允许重复采样 ) # 创建DataLoader dataloader = DataLoader( dataset, batch_size=32, sampler=sampler, # 使用我们的采样器 num_workers=4 )

3.4 验证采样效果

为了验证我们的采样器是否有效，可以检查一个batch中的类别分布：

for images, labels in dataloader: print(f"当前batch中猫的数量: {(labels == 0).sum().item()}") print(f"当前batch中狗的数量: {(labels == 1).sum().item()}") break

理想情况下，猫和狗的数量应该接近1:1，而不是原始数据中的5:1。

4. 高级技巧与注意事项

4.1 处理极端不平衡数据

当数据极度不平衡时（如1:100），可以考虑以下策略：

• 结合过采样技术（如SMOTE）
• 使用分层采样（Stratified Sampling）
• 调整权重计算公式，如使用平方根或对数变换

改进的权重计算公式：

# 使用平方根平滑极端权重差异 class_weights = [1./math.sqrt(cat_count), 1./math.sqrt(dog_count)]

4.2 与损失函数权重结合

为了获得更好的效果，可以将采样策略与损失函数权重结合：

# 在损失函数中设置类别权重 criterion = torch.nn.CrossEntropyLoss( weight=torch.tensor([1., 5.]) # 给予狗更高的惩罚权重 )

4.3 验证集和测试集的注意事项

• 验证集：应保持原始分布以反映真实场景
• 测试集：绝对不要使用采样器，保持原始分布

# 验证集DataLoader不应使用采样器 val_dataloader = DataLoader( val_dataset, batch_size=32, shuffle=False, # 验证集通常不shuffle num_workers=4 )

4.4 可视化采样效果

使用matplotlib可视化采样前后的类别分布：

import matplotlib.pyplot as plt # 原始分布 plt.bar(['猫', '狗'], [cat_count, dog_count]) plt.title('原始数据分布') plt.show() # 采样后分布（统计100个batch） sampled_cats = 0 sampled_dogs = 0 for i, (_, labels) in enumerate(dataloader): sampled_cats += (labels == 0).sum().item() sampled_dogs += (labels == 1).sum().item() if i >= 100: break plt.bar(['猫', '狗'], [sampled_cats, sampled_dogs]) plt.title('采样后数据分布(100个batch)') plt.show()

5. 替代方案比较

WeightedRandomSampler并非唯一解决方案，下表比较了几种常见方法：

在实际项目中，我通常会先尝试WeightedRandomSampler，因为它：

1. 不需要修改模型结构
2. 计算开销小
3. 效果立竿见影

但对于极端不平衡的数据（如欺诈检测中的1:10000比例），可能需要结合过采样和特殊损失函数。