Train stratifié / Test-split dans scikit-learn


88

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.

Réponses:


153

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


1
@AndreasMueller Existe-t-il un moyen simple de stratifier les données de régression?
Jordan le

3
@Jordan rien n'est implémenté dans scikit-learn. Je ne connais pas de méthode standard. Nous pourrions utiliser des centiles.
Andreas Mueller le

@AndreasMueller Avez-vous déjà vu le comportement où cette méthode est considérablement plus lente que le StratifiedShuffleSplit? J'utilisais l'ensemble de données MNIST.
snymkpr

@activatedgeek cela semble très étrange, car train_test_split (... stratify =) appelle juste StratifiedShuffleSplit et prend le premier split. N'hésitez pas à ouvrir un problème sur le tracker avec un exemple reproductible.
Andreas Mueller

@AndreasMueller En fait, je n'ai pas ouvert de problème parce que j'ai le sentiment fort que je fais quelque chose de mal (même si ce n'est que 2 lignes). Mais si je suis encore capable de le reproduire plusieurs fois aujourd'hui, je le ferai!
snymkpr

29

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

5
Notez que depuis 0.18.x, n_iterdevrait être n_splitspour StratifiedShuffleSplit - et qu'il existe une API légèrement différente pour cela: scikit-learn.org/stable/modules/generated
...

2
S'il ys'agit d'une série Pandas, utilisezy.iloc[train_index], y.iloc[test_index]
Owlright

1
@Owlright J'ai essayé d'utiliser un dataframe pandas et les indices renvoyés par StratifiedShuffleSplit ne sont pas les indices du dataframe. dataframe index: 2,3,5 the first split in sss:[(array([2, 1]), array([0]))]:(
Meghna Natraj

2
@tangy pourquoi est-ce une boucle for? n'est-il pas vrai que lorsqu'une ligne X_train, X_test = X[train_index], X[test_index]est invoquée, elle remplace X_trainet X_test? Pourquoi alors pas un seul next(sss)?
Bartek Wójcik

13

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
)
  • 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.

12

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


6

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

0
#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)

0

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]
En utilisant notre site, vous reconnaissez avoir lu et compris notre politique liée aux cookies et notre politique de confidentialité.
Licensed under cc by-sa 3.0 with attribution required.