PyTorch 不均衡数据集的采样策略

引言

在许多机器学习任务中，我们常常面临着某些类别的样本数量远远超过其他类别的情况，这种不均衡的数据分布可能会对模型的性能产生负面影响。例如，在医学诊断中，正常样本数量往往远远多于罕见病样本数量。因此，我们需要确保训练的模型不会过于偏向拥有更多样本的类别，而忽略了样本稀缺的类别。

类别不均衡问题引入

假设我们有一个罕见病诊断的数据集，其中包含了 25 张图像，其中 5 个为罕见病样本，剩余的 20 个为正常样本。

如果我们的模型简单地预测所有图像都属于正常类别，那么根据此模型计算的准确率和召回率分别为 80% 和 100%。这表明了即使模型仅仅预测了一个类别，也能够获得相当高的准确率和召回率，但这种情况下，模型可能会偏向预测多数类别，而忽略了少数类别。

数据重采样

为了解决数据不均衡的问题，我们可以使用数据重采样技术，其中包括下采样和过采样。下采样是指从多数类别的样本中随机删除一些样本，以平衡不同类别的数量；而过采样则是向少数类别的样本中添加更多的样本，从而达到类别平衡的目的。然而，这两种方法都存在一些缺点，比如过采样可能导致模型过拟合，而下采样可能导致信息丢失。

不均衡数据集采样器

为了解决上述问题，我们提出了一个名为 ImbalancedDatasetSampler 的 PyTorch 采样器。该采样器可以从不均衡的数据集中重新平衡类别分布，并自动计算采样时的权重。它的优点包括：

• 能够有效地从不平衡的数据集中采样，并保持类别之间的平衡；
• 自动估计采样时的权重，无需手动设置；
• 无需创建新的平衡数据集，节省内存和计算资源；
• 当与数据增强技术一起使用时，可以减轻过拟合的风险。

使用方法

安装 ImbalancedDatasetSampler 后，只需将其作为 DataLoader 的采样器参数即可。例如：

<code>from torchsampler import ImbalancedDatasetSampler train_loader = torch.utils.data.DataLoader( train_dataset, sampler=ImbalancedDatasetSampler(train_dataset), batch_size=args.batch_size, **kwargs )</code>

性能验证

我们以不均衡手写字符分类数据集为例，对比了使用普通采样器和 ImbalancedDatasetSampler 训练模型的性能。结果显示，在使用 ImbalancedDatasetSampler 时，少数类别的准确率明显提高，同时维持了其他类别的准确率。

结论

ImbalancedDatasetSampler 为解决不均衡数据集问题提供了一个简单而有效的解决方案。它可以帮助我们更好地训练模型，提高少数类别的识别率，同时减少过拟合的风险。因此，在实际应用中，我们强烈推荐使用该采样器来处理不均衡的数据集。