sklearn 中每个班级的具体测试数量/训练规模

Data:

import pandas as pd
data = 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]})

My code:

import numpy as np
from sklearn.cross_validation import train_test_split
X = 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将从所有类别中随机选择测试集。有什么办法可以拥有the same测试集的数量每堂课？ （例如，第1类的两个数据和第2类的两个数据。注意每个类的总数不相等）

预期结果：

y_test
array([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. 执行 groupbyX和样品n每个组的值。
2. 获取该对象的内部索引。这是我们测试集的索引，它与原始数据的集合差就是我们的训练索引。