Train stratifié / Test-split dans scikit-learn

Je dois diviser mes données en un ensemble d'entraînement (75%) et un ensemble de test (25%). Je le fais actuellement avec le code ci-dessous:

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

Cependant, j'aimerais stratifier mon ensemble de données d'entraînement. Comment je fais ça? J'ai étudié la StratifiedKFoldméthode, mais ne me laisse pas spécifier la répartition 75% / 25% et stratifier uniquement l'ensemble de données d'entraînement.

[mise à jour pour 0.17]

Voir la documentation de 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 pour 0.17]

Il y a une pull request ici . Mais vous pouvez simplement faire train, test = next(iter(StratifiedKFold(...))) et utiliser les indices de train et de test si vous le souhaitez.

TL; DR: utilisez StratifiedShuffleSplit avectest_size=0.25

Scikit-learn fournit deux modules pour le fractionnement stratifié:

1. StratifiedKFold : Ce module est utile en tant qu'opérateur de validation croisée directe de K-Fold: car il mettra en place des ensembles de n_foldsformation / test de sorte que les classes soient également équilibrées dans les deux.

Voici du code (directement à partir de la documentation ci-dessus)

>>> 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 : ce module crée un seul ensemble d'entraînement / de test ayant des classes (stratifiées) également équilibrées. C'est essentiellement ce que vous voulez avec le n_iter=1. Vous pouvez mentionner la taille du test ici comme danstrain_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

Voici un exemple de données continues / de régression (jusqu'à ce que ce problème sur GitHub soit résolu).

min = np.amin(y)
max = np.amax(y)

# 5 bins may be too few for larger datasets.
bins     = np.linspace(start=min, stop=max, num=5)
y_binned = np.digitize(y, bins, right=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, 
    y, 
    stratify=y_binned
)

• Où startest min et stopmax de votre cible continue.
• Si vous ne définissez pas, right=Truecela fera plus ou moins de votre valeur maximale un bac séparé et votre fractionnement échouera toujours car trop peu d'échantillons seront dans ce bac supplémentaire.

Vous pouvez simplement le faire avec la train_test_split()méthode disponible dans Scikit learn:

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

J'ai également préparé un court GitHub Gist qui montre comment l' stratifyoption fonctionne:

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

En plus de la réponse acceptée par @Andreas Mueller, je veux simplement ajouter cela comme @tangy mentionné ci-dessus:

StratifiedShuffleSplit ressemble le plus à train_test_split ( stratify = y) avec des fonctionnalités supplémentaires de:

1. stratifier par défaut
2. en spécifiant n_splits , il divise à plusieurs reprises les données

#train_size is 1 - tst_size - vld_size
tst_size=0.15
vld_size=0.15

X_train_test, X_valid, y_train_test, y_valid = train_test_split(df.drop(y, axis=1), df.y, test_size = vld_size, random_state=13903) 

X_train_test_V=pd.DataFrame(X_train_test)
X_valid=pd.DataFrame(X_valid)

X_train, X_test, y_train, y_test = train_test_split(X_train_test, y_train_test, test_size=tst_size, random_state=13903)

Mise à jour de la réponse @tangy ci-dessus vers la version actuelle de scikit-learn: 0.23.2 ( documentation StratifiedShuffleSplit ).

from sklearn.model_selection import StratifiedShuffleSplit

n_splits = 1  # We only want a single split in this case
sss = StratifiedShuffleSplit(n_splits=n_splits, test_size=0.25, random_state=0)

for train_index, test_index in sss.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]