如何用WeightedRandomSampler平衡PyTorch中的不平衡数据？

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

# 假设我们有一个简单的数据集
class SimpleDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 示例数据
data = torch.randn(100, 3)  # 100个样本，每个样本3个特征
labels = torch.randint(0, 2, (100,))  # 二分类标签

# 计算每个类别的权重
class_counts = torch.bincount(labels)
weights = 1.0 / class_counts
sample_weights = weights[labels]

# 创建WeightedRandomSampler
sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=True)

# 使用DataLoader加载数据
dataloader = DataLoader(dataset=SimpleDataset(data, labels), sampler=sampler, batch_size=10)

# 验证采样结果
for batch in dataloader:
    print(batch[1].unique())  # 查看每批次的标签分布