Stratified Train/Test-split in scikit-learn

Question

I need to split my data into a training set (75%) and test set (25%). I currently do that with the code below:

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

However, I'd like to stratify my training dataset. How do I do that? I've been looking into the StratifiedKFold method, but doesn't let me specifiy the 75%/25% split and only stratify the training dataset.

Answer 1

[update for 0.17]

See the docs of 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) 

[/update for 0.17]

There is a pull request here. But you can simply do train, test = next(iter(StratifiedKFold(...))) and use the train and test indices if you want.

Answer 2

TL;DR : Use StratifiedShuffleSplit with test_size=0.25

Scikit-learn provides two modules for Stratified Splitting:

1. StratifiedKFold : This module is useful as a direct k-fold cross-validation operator: as in it will set up n_folds training/testing sets such that classes are equally balanced in both.

Heres some code(directly from above documentation)

>>> 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 : This module creates a single training/testing set having equally balanced(stratified) classes. Essentially this is what you want with the n_iter=1. You can mention the test-size here same as in train_test_split

Code:

>>> 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