使用 WeightedRandomSampler 解决数据样本不均衡的问题

问题描述

数据集样本不均衡。

例如,一个二分类任务,标签为 0 的数据占了 90%,标签为 1 的数据却只占 10%,用全部原始数据训练模型很可能导致模型带有一定的 ”偏见“,也可能会导致模型训练效率很慢。

使用 WeightedRandomSampler 均衡数据

PyTorch 官方文档:

以下即是 PyTorch 官方文档给出的说明:

可以看到官方文档中给了代码示例:

>>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
[4, 4, 1, 4, 5]
>>> list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False))
[0, 1, 4, 3, 2]

结合这个示例解释一下各个参数:

    weights:用来采样的权重序列,并且并不要求这个序列的加和为 1。 比如第一个示例中的 [0.1, 0.9, 0.4, 0.7, 3.0, 0.6],就表示取第 0 个样本的权重(可以理解为概率,但不等同)为 0.1,取第 1 个样本的权重为 0.9,…,取第 4 个样本的权重为 3.0,…,很明显,取第 4 个样本的概率最大,然后是 第 1 个、第 3 个、第 5 个、第 2 个、第 0 个。 num_samples:采样的数量。 比如示例中 num_samples 为 5,即采样数量为 5,最终输出的序列长度也为 5。 replacement:是否可放回采样。

应用

以下是使用 WeightedRandomSampler 来解决数据集样本不均衡问题的代码。

from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler


class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, index):
        return {
          
   
            "text": self.data[index]["text"],
            "label": self.data[index]["label"]
        }

    def __len__(self):
        return len(self.data)


if __name__ == __main__:
    data = [
        {
          
   "text": "a", "label": 0}, {
          
   "text": "b", "label": 0}, {
          
   "text": "c", "label": 1}, {
          
   "text": "d", "label": 0},
        {
          
   "text": "e", "label": 0}, {
          
   "text": "f", "label": 0}, {
          
   "text": "g", "label": 0}, {
          
   "text": "h", "label": 0},
        {
          
   "text": "i", "label": 0}, {
          
   "text": "j", "label": 0}, {
          
   "text": "k", "label": 0}, {
          
   "text": "l", "label": 1}
    ]
    dataset = MyDataset(data)

    label_list = []
    for per_data in dataset:
        label_list.append(per_data["label"])
    print(f"label_list = {
            
     label_list}")

    weights = [1.0 / label_list.count(label) for label in label_list]
    print(f"weights = {
            
     weights}")

    sampler = WeightedRandomSampler(weights, len(dataset), replacement=True)
    train_loader = DataLoader(dataset, sampler=sampler, batch_size=4, shuffle=False, num_workers=0)
    for data in train_loader:
        print(data)

输出:

label_list = [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1]
weights = [0.1, 0.1, 0.5, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.5]
{
          
   text: [c, a, g, l], label: tensor([1, 0, 0, 1])}
{
          
   text: [e, j, l, k], label: tensor([0, 0, 1, 0])}
{
          
   text: [c, g, c, k], label: tensor([1, 0, 1, 0])}

数据集有 12 条数据,其中包含 10 个标签为 0 的数据和 2 个标签为 1 的,可以说数据集是及不均衡了。

首先计算权重序列,标签为 0 的权重均为 1 / 10 = 0.1 1 / 10=0.1 1/10=0.1,标签为 1 的权重均为 1 / 2 = 0.5 1 / 2=0.5 1/2=0.5,把标签为 0 的权重全加起来的值和把标签为 1 的权重全加起来的值相等的,因此采样时随机取到标签为 0 和标签为 1 的概率是相等的,最终新的 dataloader 中基本上数据就均衡了,不过因为是随机采样,两种标签的数据量最终差一、两个也是可以理解的,总之是比原始数据集均衡了很多。

经验分享 程序员 微信小程序 职场和发展