PyTorch中WeightedRandomSampler处理类别不平衡的原理解惑
你当前的计算逻辑是对的,但存在一个容易出错的细节问题:
PyTorch张量的value_counts()方法默认会按样本数量降序返回统计结果,返回的顺序和你原始的类别索引(0、1、2、3)不匹配。你当前得到的class_sample_count = [2555, 2552, 621, 227]对应的类别依次是「类别0、类别3、类别2、类别1」,直接用weight[train_labels]取值时,会出现类别索引和权重不匹配的错误,比如标签为1的样本会匹配到类别3的权重。
修正方案如下,按类别索引0~3的顺序统计样本数即可:
class_sample_count = np.array([len(train_labels[train_labels == i]) for i in range(4)])
修正后class_sample_count的顺序会和你的原始类别一一对应:[2555, 227, 621, 2552],后续计算的权重才会正确匹配到对应类别的样本。
你看到的「张量打印出来所有样本权重近似相同」是PyTorch默认打印精度的问题,不是权重真的没有差异:PyTorch默认打印浮点型张量时只会保留4位小数,你自己打印的原始数值里,类别0的权重约为0.000391,类别1的权重约为0.004405,两者实际差了11倍,只是四舍五入后都显示为0.0004而已。
WeightedRandomSampler的工作逻辑是按每个样本的权重占所有样本总权重的比例来进行随机采样,和权重的绝对大小无关:
当你按修正后的方法计算权重后,每个类别的所有样本的权重总和是相等的:
- 类别0总权重 = 2555 * (1/2555) = 1
- 类别1总权重 = 227 * (1/227) = 1
- 类别2总权重 = 621 * (1/621) = 1
- 类别3总权重 = 2552 * (1/2552) = 1
四个类别的总权重占比均为25%,所以采样时每个类别被抽到的概率完全相等,刚好可以解决原始数据集的类别不平衡问题。
内容的提问来源于stack exchange,提问作者sidra Aleem

