如何在Scikit-learn中使用train_test_split控制类别比例划分数据

这问题我太熟了！你遇到的是不平衡数据集划分时类别比例失控的典型情况——因为默认的train_test_split是完全随机抽样，对于像b这种只占10%的小样本类别，很容易出现抽不到或者比例严重偏离原数据的情况。

核心解决方案：使用stratify参数 #

Scikit-learn的train_test_split专门提供了stratify参数，它会严格按照原始数据的类别比例来划分训练集和测试集，确保两边的类别分布和原数据保持一致。

修改你的代码如下：

from sklearn.model_selection import train_test_split
from collections import Counter
x_coords = range(100)
labels = ['a']*90 + ['b']*10
# 新增stratify参数，传入标签列表
x_train, x_test, label_train, label_test = train_test_split(x_coords, labels, test_size=20, stratify=labels)
print(Counter(label_test))

效果说明 #

修改后每次运行，测试集的类别比例都会严格遵循原数据的9:1，也就是稳定输出Counter({'a': 18, 'b': 2})，再也不会出现全是a类的情况。

另外这个参数对多分类问题同样适用，不管你有多少个类别，它都会保证每个类别的划分比例和原数据一致。

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