web-dev-qa-db-fra.com

Train stratifié/Test divisé en scikit-learn

Je dois diviser mes données en un ensemble de formation (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 jeu de données d'entraînement. Comment je fais ça? J'ai étudié la méthode StratifiedKFold, mais je ne me laisse pas spécifier le partage 75%/25% et seulement stratifier le jeu de données d'apprentissage.

pythonscikit-learn
44
pir

[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 demande de tir ici . Mais vous pouvez simplement faire train, test = next(iter(StratifiedKFold(...))) Et utiliser les indices train et test si vous le souhaitez.

93
Andreas Mueller

TL; DR: Utiliser StratifiedShuffleSplit avec test_size=0.25

Scikit-learn propose deux modules pour le fractionnement en couches:

  1. StratifiedKFold : Ce module est utile en tant qu'opérateur de validation croisée à pliages multiples direct: il configurera n_folds ensembles de formation/tests de sorte que les classes soient équilibrées dans les deux cas.

Heres 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 de formation/test ayant des classes également équilibrées (stratifiées). En gros, c’est ce que vous voulez avec le n_iter=1. Vous pouvez mentionner ici la même taille que dans 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
18
tangy

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

# 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)
8
Jordan

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

StratifiedShuffleSplit ressemble le plus étroitement à train_test_split (stratify = y) Avec en plus les caractéristiques suivantes:

  1. stratifier par défaut
  2. en spécifiant n_splits, il divise à plusieurs reprises les données 
4
Max
#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
José Carlos Castro

Vous pouvez simplement le faire avec la méthode train_test_split() 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 fonctionne l'option stratify:

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

0
Shayan Amani