在scikit-learn中分层训练/测试分裂

Question

我需要将我的数据分成训练集(75%)和测试集(25%).我目前使用以下代码执行此操作:

X, Xt, userInfo, userInfo_train = sklearn.cross_validation.train_test_split(X, userInfo)   

但是,我想对训练数据集进行分层.我怎么做？我一直在研究这种StratifiedKFold方法,但是不允许我指定75%/ 25%的分割,只对训练数据集进行分层.

Answer 1

[更新0.17]

查看以下文档sklearn.model_selection.train_test_split:

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    stratify=y, 
                                                    test_size=0.25)

[/更新0.17]

有一个拉要求在这里.但train, test = next(iter(StratifiedKFold(...))) 如果你愿意,你可以简单地使用火车和测试指数.

Answer 2

TL; DR:使用StratifiedShuffleSplit与test_size=0.25

Scikit-learn为分层分裂提供了两个模块:

1. StratifiedKFold:此模块可用作直接k折叠交叉验证运算符:因为它将设置n_folds训练/测试集,使得类在两者中均衡.

下面是一些代码(直接来自上面的文档)

>>> skf = cross_validation.StratifiedKFold(y, n_folds=2) #2-fold cross validation
>>> len(skf)
2
>>> for train_index, test_index in skf:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
...    #fit and predict with X_train/test. Use accuracy metrics to check validation performance

2. StratifiedShuffleSplit:该模块创建一个具有同等平衡(分层)类的训练/测试集.基本上这就是你想要的n_iter=1.您可以在此处提及与测试大小相同的测试大小train_test_split

码:

>>> sss = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=0)
>>> len(sss)
1
>>> for train_index, test_index in sss:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
>>> # fit and predict with your classifier using the above X/y train/test

Answer 3

这是连续/回归数据的示例(直到GitHub上的此问题得到解决).

# Your bins need to be appropriate for your output values
# e.g. 0 to 50 with 25 bins
bins     = np.linspace(0, 50, 25)
y_binned = np.digitize(y_full, bins)
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y_binned)

Answer 4

您可以train_test_split()使用Scikit学习中可用的方法简单地做到这一点：

from sklearn.model_selection import train_test_split 
train, test = train_test_split(X, test_size=0.25, stratify=X['YOUR_COLUMN_LABEL']) 

我还准备了一个简短的GitHub Gist，展示了stratifyoption的工作原理：

https://gist.github.com/SHi-ON/63839f3a3647051a180cb03af0f7d0d9

Answer 5

除了@Andreas Mueller接受的答案外，只需将其添加为上述@tangy：

StratifiedShuffleSplit最类似于train_test_split（stratify = y），具有以下新增功能：

1. 默认分层
2. 通过指定n_splits，它反复拆分数据