WeightedRandomSampler 详解及示例分析

应用场景：

我们常常会遇到--数据集不平衡--问题，举个最简单的例子：猫狗分类，你有10000张狗的图片，却只有1000张猫的图片。这时如果直接利用整个数据集训练就很容易导致网络对‘狗’这个类别过拟合，在猫的识别任务上表现很差。

解决方法：使用pytorch的WeightedRandomSampler，每个epoch根据每个类别的数量赋予权重，对每个类别样本进行采样。

例如：每个epoch使用1600张图片训练网络，其中狗的样本被采样的概率低，但数量多，猫的样本被采样概率高，但数量少，这就使得1600张训练图片中两个类别数量差不多，而且每个epoch都进行采样，训练够多的话，所有图片基本都被训练过。

---

官方用法介绍

先看下官网的API

weight：一个列表，其中每个数衡量位于该index的样本被采样的概率权重（不需要总和为1）
num_samples：一个整数，表示你要采样的个数，比如上面猫狗分类每epoch在整个数据集中采样1600个图片
replacement：一个bool变量，是否重复采样，如果是True，则一个图片可能被采样多次。
generator：没用，不用管。

---

应用示例：

猫狗分类，10000张狗图片，1000张猫图片。要实现：每个epoch只取1600张图片训练，且不重复采样

关键代码：

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

# 自定义的dataset，前10000个是dog，然后是1000个cat，共11000个数据。
train_dataset   = DataGenerator(train_imgs, input_shape=input_shape, train=True)
# 每个类别的样本数量
num_dog   	    = 10000
num_cat         = 1000
# weight列表
train_weights   = []
train_weights.extend([1/num_dog]*num_dog) # 扩展10000个0.0001，前10000个样本的采样权重，狗多权重小
train_weights.extend([1/num_cat]*num_cat)  # 扩展1000个0.001，后1000个样本的采样权重，猫少权重大
print(train_weights) # 可以打印看看，总共11000个数，前10000个是0.0001,然后是1000个0.001
# 创建WeightedRandomSampler，1600为采样数
train_sampler   = WeightedRandomSampler(train_weights, 1600， replacement=False)
# 将sampler传给Dataloader，不再需要shuffle
gen             = DataLoader(train_dataset, batch_size=128, sampler=train_sampler, num_workers=8)

解释：weight列表里面是每个样本的采样概率，11000个样本那就有11000个权重，例子中，狗的权重加起来为1，猫的权重加起来也是1，猫狗权重总和相等，保证每个epoch中1600样本中两类图片样本数差不多。Dataloader每次先利用sampler从dataset中采样1600个图片，再分成一堆堆的batch进行训练。

怕大家看不懂，废话很多请见谅。如果有问题请评论区交流。