Python中WeightedRandomSampler()的简单使用教程

WeightedRandomSampler是PyTorch中的一个采样器（Sampler），它可以根据样本的权重来对样本进行采样。在某些情况下，我们希望在训练过程中更多地关注一些少数类别的样本，以提高模型对这些类别的学习效果。这时，可以使用WeightedRandomSampler来实现有偏采样，即样本的采样概率与其权重成正比。

下面是WeightedRandomSampler的简单使用教程及示例：

首先，我们需要导入必要的库和模块：

import torch
import torchvision
from torch.utils.data import WeightedRandomSampler

接着，我们需要准备一个带有样本权重的数据集。以FashionMNIST数据集为例，假设我们希望在训练过程中更关注类别为8（鞋子）和9（包）的样本，可以创建一个权重向量weights，其中类别8和9的权重分别设为2和3，其他类别的权重都设为1：

# 加载FashionMNIST数据集
train_dataset = torchvision.datasets.FashionMNIST(root='data/',
                                                 train=True,
                                                 transform=torchvision.transforms.ToTensor(),
                                                 download=True)
# 获取训练集样本数量
num_train = len(train_dataset)

# 设置样本权重
weights = torch.zeros(num_train)
targets = train_dataset.targets

# 设定类别8和9的样本权重为2和3，其他类别的权重为1
weights[targets == 8] = 2
weights[targets == 9] = 3

然后，我们可以使用WeightedRandomSampler来创建一个采样器，并将其与数据集一起传递给DataLoader：

# 创建一个采样器
sampler = WeightedRandomSampler(weights, num_samples=num_train, replacement=True)

# 创建一个DataLoader，并将采样器与数据集一起传递
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=64,
                                           sampler=sampler)

在上述代码中，通过指定参数weights来传递样本权重，num_samples参数设置为num_train即样本的总数，replacement参数设为True表示采样时可以有重复的样本。

最后，我们可以使用train_loader来迭代训练数据集：

for images, labels in train_loader:
    # 在这里进行训练操作
    pass

以上就是WeightedRandomSampler的简单使用教程和示例。通过设置样本的权重，我们可以实现有偏采样，以更关注某些少数类别的样本。这在处理类别不平衡的数据集时非常有用，可以提高模型对少数类别的学习效果。