简介

在许多机器学习应用中，我们经常会遇到一些数据集中，某些类别的样本数量远多于其他类别。以罕见病检测为例，正常样本的数量通常会远远超过患病样本。在这种情况下，我们需要确保训练出的模型不会偏向于样本数量较多的类别。例如，假设一个数据集中有5张患病图像和20张正常图像。如果模型将所有图像都预测为正常，那么它的准确率将达到80%，F1分数也会达到0.88。然而，这样的模型显然存在严重的偏差，倾向于将所有样本分类为“正常”。

为了解决这一问题，一种广泛采用的技术称为重采样。它包括从多数类中移除样本（欠采样）和/或从少数类中添加更多样本（过采样）。尽管这些技术能够有效平衡各类别之间的分布，但它们也各有其局限性（天下没有免费的午餐）。最简单的过采样方法是随机复制少数类中的样本，但这可能导致过拟合。而最简单的欠采样方法则是随机删除多数类中的样本，这可能会导致信息丢失。

在这个仓库中，我们实现了一个易于使用的 PyTorch 采样器 ImbalancedDatasetSampler，它能够：

• 在从不平衡数据集中采样时重新平衡各类别分布
• 自动估计采样权重
• 避免创建一个新的平衡数据集
• 结合数据增强技术使用时，可以有效缓解过拟合问题

使用方法

要快速开始使用，可以通过以下方式安装该包：

pip install torchsampler

在创建 DataLoader 时，只需将 ImbalancedDatasetSampler 作为 sampler 参数传入即可。例如：

from torchsampler import ImbalancedDatasetSampler

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    sampler=ImbalancedDatasetSampler(train_dataset),
    batch_size=args.batch_size,
    **kwargs
)

这样，在每个 epoch 中，数据加载器都会对整个数据集进行采样，并根据每个类别的出现概率反向调整样本权重。

示例：不平衡 MNIST 数据集

不平衡数据集中各类别的分布情况如下：

使用不平衡数据集采样器后：

（左：每个 epoch 的测试准确率；右：混淆矩阵）

未使用不平衡数据集采样器时：

（左：每个 epoch 的测试准确率；右：混淆矩阵）

可以看到，对于 2、6 和 9 等小类别的分类性能有了显著提升，而其他类别的准确率则保持不变。

贡献

我们非常欢迎各种形式的贡献。如果您计划提交 bug 修复，请直接提交，无需额外讨论。如果您打算贡献新功能、实用工具函数或扩展，请先打开一个 issue，与我们讨论该功能的具体实现。

许可证