数据:

import pandas as pddata = pd.DataFrame({'classes':[1,1,1,2,2,2,2],'b':[3,4,5,6,7,8,9], 'c':[10,11,12,13,14,15,16]})

我的代码:

import numpy as npfrom sklearn.cross_validation import train_test_splitX = np.array(data[['b','c']])  y = np.array(data['classes'])     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=4)

问题:

train_test_split 会从所有类别中随机选择测试集。有没有办法让每个类别的测试集数量相同？（例如，从类别1中取两个数据，从类别2中取两个数据。请注意，每个类别的总数据量是不相等的）

预期结果:

y_testarray([1, 2, 2, 1], dtype=int64)

回答：

实际上，sklearn 中没有直接实现这一功能的函数或参数。stratify 参数是按比例抽样的，这并不是你想要的效果，如你在评论中所指出的那样。

你可以构建一个自定义函数，虽然速度相对较慢，但在绝对意义上并不算非常慢。请注意，这是为 pandas 对象设计的。

def train_test_eq_split(X, y, n_per_class, random_state=None):    if random_state:        np.random.seed(random_state)    sampled = X.groupby(y, sort=False).apply(        lambda frame: frame.sample(n_per_class))    mask = sampled.index.get_level_values(1)    X_train = X.drop(mask)    X_test = X.loc[mask]    y_train = y.drop(mask)    y_test = y.loc[mask]    return X_train, X_test, y_train, y_test

示例案例:

data = pd.DataFrame({'classes': np.repeat([1, 2, 3], [10, 20, 30]),                     'b': np.random.randn(60),                     'c': np.random.randn(60)})y = data.pop('classes')X_train, X_test, y_train, y_test = train_test_eq_split(    data, y, n_per_class=5, random_state=123)y_test.value_counts()# 3    5# 2    5# 1    5# Name: classes, dtype: int64

工作原理:

1. 对 X 进行分组，并从每个组中抽取 n 个值。
2. 获取该对象的内层索引。这就是我们的测试集索引，其与原始数据的集合差就是我们的训练集索引。