python sklearn的训练集测试集根据标签的比例来切分数据集

在切分数据集的时候发现标签样本非常不均衡的情况, 这时候切分数据比为了避免没有均衡切分每一种分类就可以用sklearn来实现

train_x,test_x,train_y,test_y = train_test_split(x_data, y_data, test_size = 0.2,random_state = 32, stratify=y_data)

其中stratify默认是none 但是如果要根据标签(y) 来进行按比例分割就可以加上stratify = y_data (y数据)

然后分完数据集如果想要查看是否按比例分配:

可以计算list中每个元素出现的次数 

from collections import Counter
a = Counter(train_y) # 训练集里面各个标签的出现次数
b = Counter(test_y) # 测试集里面各个标签的出现次数
c = Counter(labels) # 未切分前各个标签的出现次数
print(a)
print(b)
print(c)

我的结果如下: