sklearn

Question

数据：

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]})

我的代码：

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将从所有类中随机选择测试集。是否有任何方法可以为每个类设置相同的测试集数量？ （例如，来自类1的两个数据和来自类2的两个数据。请注意，每个类的总数不相等）

预期结果：

y_test
array([1, 2, 2, 1], dtype=int64)

Answer 1

实际上没有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. 获取此对象的内部索引。这是我们的测试集的索引，它与原始数据的设置差异是我们的列车索引。