在scikit-learn中进行分层训练/验证/测试拆分 [英] Stratified Train/Validation/Test-split in scikit-learn

问题描述

这里已经有关于如何通过train_test_split(如何将数据分成3组(训练，验证和测试)?).但是如何进行分层训练/验证/测试拆分呢?

There is already a description here of how to do stratified train/test split in scikit via train_test_split (Stratified Train/Test-split in scikit-learn) and a description of how to random train/validation/test split via np.split (How to split data into 3 sets (train, validation and test)?). But what about doing stratified train/validation/test split.

进行分层(在类标签上)训练/验证/测试拆分时想到的最接近的近似值如下，但我怀疑有更好的方法也许可以通过一个函数调用或更准确地实现:

The closest approximation that comes to mind for doing stratified (on class label) train/validation/test split is as follows, but I suspect there's a better way that can perhaps achieve this in one function call or in a more accurate way:

假设我们要进行60/20/20训练/验证/测试拆分，那么我目前的方法是先对60/40进行分层拆分，然后对前40进行50/50分层拆分.最终得到60/20/20的分层拆分.

Let's say we want to do a 60/20/20 train/validation/test split, then my current approach is to first do 60/40 stratified split, then do a 50/50 stratifeid split on that first 40 as to ultimately get a 60/20/20 stratified split.

from sklearn.cross_validation import train_test_split
SEED = 2000
x_train, x_validation_and_test, y_train, y_validation_and_test = train_test_split(x, y, test_size=.4, random_state=SEED)
x_validation, x_test, y_validation, y_test = train_test_split(x_validation_and_test, y_validation_and_test, test_size=.5, random_state=SEED)

如果我的方法正确和/或您有更好的方法，请回来.

Please get back if my approach is correct and/or if you have a better approach.

谢谢

推荐答案

解决方案是仅使用StratifiedShuffleSplit两次，如下所示:

The solution is to just use StratifiedShuffleSplit twice, like below:

from sklearn.model_selection import StratifiedShuffleSplit

split = StratifiedShuffleSplit(n_splits=1, test_size=0.4, random_state=42)
for train_index, test_valid_index in split.split(df, df.target):
    train_set = df.iloc[train_index]
    test_valid_set = df.iloc[test_valid_index]

split2 = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=42)
for test_index, valid_index in split2.split(test_valid_set, test_valid_set.target):
    test_set = test_valid_set.iloc[test_index]
    valid_set = test_valid_set.iloc[valid_index]

这篇关于在scikit-learn中进行分层训练/验证/测试拆分的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！