平衡多分类数据集拆分后测试集失衡的原因及是否需平衡咨询

问题原因 #

你使用的train_test_split默认采用纯随机抽样（stratify参数默认值为None），不会按照原始数据集的类别比例进行拆分。虽然你的原始数据集是8类各100个样本的平衡分布，但单类样本数（100个）不算特别多，按20%比例抽取测试集时，随机抽样的概率波动会导致部分类别在测试集中的数量偏离预期的20个（100*0.2），出现第二类仅10个样本的情况是概率上的正常现象。

如果想要拆分后训练集和测试集都严格保持原始类别比例，只需在train_test_split中指定stratify参数为标签列即可启用分层抽样，修正后的代码如下：

X_train, X_test, y_train, y_test = train_test_split(data.iloc[:, :-1], data.iloc[:, -1], test_size=0.2, random_state=42, stratify=data.iloc[:, -1])

是否需要平衡测试集？ #

绝对不需要人为调整测试集的类别分布。

测试集的核心作用是模拟模型在真实业务场景中的数据分布，强行平衡测试集会让模型评估结果脱离实际：比如真实场景中某类样本本来就稀缺，人为补全后，模型的准确率、召回率等指标会被高估，无法反映模型在真实环境下的表现。

如果你的任务关注少数类的识别效果，正确的做法是在训练阶段处理：比如对训练集的少数类做过采样、对多数类做欠采样，或者在模型训练时使用加权损失函数（给少数类样本更高的权重）。测试集需保持拆分后的原始状态，同时改用更适合不平衡数据的评估指标（比如F1分数、混淆矩阵、AUC-ROC，而非仅看整体准确率）来衡量模型性能。

内容的提问来源于stack exchange，提问作者kkk123