Pytorch 类别平衡化处理

采用  WeightedRandomSampler:

def make_weights_for_balanced_classes(images, nclasses):

    count = [0] * nclasses

    for item in images:

        count[item[1]] += 1

    weight_per_class = [0.] * nclasses

    N = float(sum(count))

    for i in range(nclasses):

        weight_per_class[i] = N/float(count[i])

    weight = [0] * len(images)

    for idx, val in enumerate(images):

        weight[idx] = weight_per_class[val[1]]

    return weight      

dataset_train = datasets.ImageFolder(traindir)                                                                         

# For unbalanced dataset we create a weighted sampler

weights = make_weights_for_balanced_classes(dataset_train.imgs, len(dataset_train.classes))

weights = torch.DoubleTensor(weights)

sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, len(weights))                     

train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=args.batch_size, shuffle = True,

                                                             sampler = sampler, num_workers=args.workers, pin_memory=True)     

Reference:   Balanced Sampling between classes with torchvision DataLoader

参考方法2: 作者给出了均匀采样和非均匀采样的差别

imbalanced-dataset-sampler