用PyTorch的WeightedRandomSampler解决数据不平衡问题：从原理到实战

当你面对一个猫狗分类数据集，发现其中90%的样本都是狗，只有10%是猫时，传统的随机采样方法会让模型严重偏向于识别狗。这种数据不平衡问题在医疗影像分析、金融欺诈检测等领域尤为常见。本文将带你深入理解PyTorch中的WeightedRandomSampler，从底层原理到完整代码实现，彻底解决少数类样本被模型忽视的问题。

1. 数据不平衡问题的本质与影响

数据不平衡是指不同类别的样本数量存在显著差异。以信用卡欺诈检测为例，正常交易可能占99.9%，而欺诈交易仅占0.1%。这种极端不平衡会导致模型倾向于预测多数类，因为即使完全忽略少数类，也能获得很高的准确率。

数据不平衡带来的主要问题包括：

• 模型偏见：模型倾向于预测多数类，忽视少数类
• 评估指标失真：准确率等传统指标无法反映真实性能
• 资源浪费：少数类样本的信息未被充分利用

常见解决方案对比：

2. WeightedRandomSampler原理解析

WeightedRandomSampler是PyTorch提供的一种采样策略，它通过为每个样本分配不同的采样权重，使得少数类样本有更高的概率被选中。其核心参数包括：

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

关键点解析：

1. 权重计算：权重应与类别频率成反比，即少数类样本获得更高权重
2. 采样过程：根据权重进行有放回或无放回抽样
3. 与DataLoader集成：作为sampler参数传递给DataLoader

注意：权重是赋给每个样本而非类别，需要为数据集中的每个样本单独计算权重

3. 实战：从数据准备到模型训练

让我们通过一个猫狗分类的完整示例，演示如何应用WeightedRandomSampler。

3.1 数据集准备与权重计算

首先，我们需要计算每个样本的权重。假设我们的数据集分布如下：

python复制# 样本标签分布示例
labels = ['cat', 'cat', 'dog', 'dog', 'dog', 'dog', 'dog']

计算权重的完整代码：

python复制import numpy as np
from collections import Counter

def calculate_sample_weights(labels):
    # 统计每个类别的出现次数
    class_counts = Counter(labels)
    
    # 计算每个类别的权重（与频率成反比）
    class_weights = {cls: 1.0/count for cls, count in class_counts.items()}
    
    # 为每个样本分配对应类别的权重
    sample_weights = [class_weights[cls] for cls in labels]
    
    return sample_weights

# 示例使用
labels = ['cat', 'cat', 'dog', 'dog', 'dog', 'dog', 'dog']
weights = calculate_sample_weights(labels)
print(f"样本权重: {weights}")

输出结果：

code复制样本权重: [0.5, 0.5, 0.2, 0.2, 0.2, 0.2, 0.2]

3.2 构建自定义数据集与采样器

接下来，我们将权重应用到PyTorch的数据加载流程中：

python复制import torch
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler

class CatDogDataset(Dataset):
    def __init__(self, images, labels):
        self.images = images
        self.labels = labels
        
    def __len__(self):
        return len(self.labels)
        
    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]

# 假设我们已经准备好了图像数据和标签
images = [...]  # 实际的图像数据
labels = [...]  # 实际的标签数据

# 创建数据集
dataset = CatDogDataset(images, labels)

# 计算样本权重
weights = calculate_sample_weights(labels)
weights = torch.DoubleTensor(weights)

# 创建采样器
sampler = WeightedRandomSampler(
    weights,
    num_samples=len(weights),  # 通常与数据集大小相同
    replacement=True           # 允许重复采样
)

# 创建DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,
    num_workers=4
)

3.3 模型训练中的注意事项

在训练循环中，使用带采样器的DataLoader与常规方式无异：

python复制model = ...  # 你的模型定义
criterion = ...  # 损失函数
optimizer = ...  # 优化器

for epoch in range(num_epochs):
    for batch_images, batch_labels in dataloader:
        # 前向传播
        outputs = model(batch_images)
        
        # 计算损失
        loss = criterion(outputs, batch_labels)
        
        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

提示：即使使用了WeightedRandomSampler，在计算损失时也可以额外应用类别权重，双重保障模型对少数类的关注

4. 高级技巧与常见问题解决

4.1 处理极端不平衡数据

对于极度不平衡的数据（如1:1000），可以考虑以下策略：

• 组合采样：同时使用过采样和欠采样
• 分层采样：确保每个batch中都包含所有类别的样本
• 动态权重调整：根据训练过程中的表现动态调整采样权重

4.2 采样器性能优化

当数据集非常大时，WeightedRandomSampler可能会成为性能瓶颈。可以考虑：

• 预计算索引：提前计算好采样索引序列
• 批量采样：每次采样一个batch而非单个样本
• 分布式采样：在多GPU训练中确保采样的一致性

4.3 验证与测试集的采样策略

需要注意的是：

• 验证集：通常保持原始分布，以评估模型在真实场景中的表现
• 测试集：绝对不要使用任何采样策略，保持完全原始分布

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

5. 替代方案与组合策略

除了WeightedRandomSampler，还有其他应对数据不平衡的方法：

损失函数层面的解决方案：

python复制# 使用带权重的交叉熵损失
class_weights = torch.tensor([1.0, 5.0])  # 假设类别1是少数类
criterion = torch.nn.CrossEntropyLoss(weight=class_weights)

数据增强技术：

• 对少数类样本应用更激进的数据增强
• 使用SMOTE等过采样技术生成合成样本

模型架构调整：

• 在最后一层添加类别偏置项
• 使用多任务学习，将少数类识别作为辅助任务

在实际项目中，我通常会尝试组合多种策略。例如，同时使用WeightedRandomSampler和带权重的损失函数，并在少数类样本上应用数据增强。这种组合方法在医疗影像分析任务中取得了显著效果，将少数类的召回率从15%提升到了68%。