如何在PyTourch中创建平衡循环迭代器？

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

# 假设有一个数据集dataset，其中包含了样本和对应的标签
dataset = ...

# 计算每个类别的样本数量
class_counts = [10, 20, 30]  # 假设有3个类别，分别有10、20、30个样本

# 计算每个样本的权重，使得每个类别的样本被平衡地采样
weights = [1.0 / count for count in class_counts]
sample_weights = [weights[label] for _, label in dataset]

# 创建平衡循环迭代器
sampler = WeightedRandomSampler(sample_weights, len(dataset), replacement=True)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

# 使用平衡循环迭代器进行训练
for inputs, labels in dataloader:
    # 在这里进行模型训练
    ...