web-dev-qa-db-fra.com

sklearn train_test_split on pandas stratifier par plusieurs colonnes

Je suis un utilisateur relativement nouveau de sklearn et j'ai rencontré un comportement inattendu dans train_test_split de sklearn.model_selection. J'ai un pandas dataframe que je voudrais diviser en un ensemble de formation et de test. Je voudrais stratifier mes données par au moins 2, mais idéalement 4 colonnes dans mon dataframe.

Sklearn n'a émis aucun avertissement lorsque j'ai essayé de le faire, mais j'ai découvert plus tard qu'il y avait des lignes répétées dans mon jeu de données final. J'ai créé un exemple de test pour montrer ce comportement:

from sklearn.model_selection import train_test_split
a = np.array([i for i in range(1000000)])
b = [i%10 for i in a]
c = [i%5 for i in a]
df = pd.DataFrame({'a':a, 'b':b, 'c':c})

Cela semble fonctionner comme prévu si je stratifie par l'une ou l'autre colonne:

train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['b']])
print(len(train.a.values))  # prints 800000
print(len(set(train.a.values)))  # prints 800000

train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['c']])
print(len(train.a.values))  # prints 800000
print(len(set(train.a.values)))  # prints 800000

Mais lorsque j'essaie de stratifier par les deux colonnes, j'obtiens des valeurs répétées:

train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['b', 'c']])
print(len(train.a.values))  # prints 800000
print(len(set(train.a.values)))  # prints 640000
pythonpandasscikit-learn
14
Caitlin

La raison pour laquelle vous obtenez des doublons est que train_test_split() définit finalement les strates comme ensemble unique de valeurs de tout ce que vous avez passé dans le stratify argument. Étant donné que les strates sont définies à partir de deux colonnes, une ligne de données peut représenter plus d'une strate, et donc l'échantillonnage peut choisir deux fois la même ligne car il pense qu'il s'agit d'un échantillonnage de différentes classes.

La fonction train_test_split()appelleStratifiedShuffleSplit, qui tilisenp.unique() sur y (qui est ce que vous passer via stratify). Du code source:

classes, y_indices = np.unique(y, return_inverse=True)
n_classes = classes.shape[0]

Voici un exemple de cas simplifié, une variante de l'exemple que vous avez fourni:

from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd

N = 20
a = np.arange(N)
b = np.random.choice(["foo","bar"], size=N)
c = np.random.choice(["y","z"], size=N)
df = pd.DataFrame({'a':a, 'b':b, 'c':c})

print(df)
     a    b  c
0    0  bar  y
1    1  foo  y
2    2  bar  z
3    3  bar  y
4    4  foo  z
5    5  bar  y
...

La fonction de stratification pense qu'il y a quatre classes sur lesquelles se diviser: foo, bar, y et z. Mais comme ces classes sont essentiellement imbriquées, ce qui signifie que y et z apparaissent toutes les deux dans b == foo et b == bar, nous obtiendrons des doublons lorsque le séparateur essaiera d'échantillonner dans chaque classe.

train, test = train_test_split(df, test_size=0.2, random_state=0, 
                               stratify=df[['b', 'c']])
print(len(train.a.values))  # 16
print(len(set(train.a.values)))  # 12

print(train)
     a    b  c
3    3  bar  y   # selecting a = 3 for b = bar*
5    5  bar  y
13  13  foo  y
4    4  foo  z
14  14  bar  z
10  10  foo  z
3    3  bar  y   # selecting a = 3 for c = y
6    6  bar  y
16  16  foo  y
18  18  bar  z
6    6  bar  y
8    8  foo  y
18  18  bar  z
7    7  bar  z
4    4  foo  z
19  19  bar  y

#* We can't be sure which row is selecting for `bar` or `y`, 
#  I'm just illustrating the idea here.

Il y a une question de conception plus large ici: voulez-vous utiliser un échantillonnage stratifié imbriqué, ou voulez-vous simplement traiter chaque classe dans df.b et df.c en tant que classe distincte pour échantillonner? Si c'est le cas, c'est ce que vous obtenez déjà. Le premier est plus compliqué, et ce n'est pas ce que train_test_split est configuré pour le faire.

Vous pourriez trouver cette discussion d'échantillonnage stratifié imbriqué utile.

11
andrew_reece

Quelle version de scikit-learn utilisez-vous? Vous pouvez utiliser sklearn.__version__ vérifier.

Avant la version 0.19.0, scikit-learn ne gère pas correctement la stratification bidimensionnelle. Il est corrigé en 0.19.0.

Il est décrit dans problème # 9044 .

Mettre à jour votre scikit-learn devrait résoudre le problème. Si vous ne pouvez pas mettre à jour votre scikit-learn, consultez cet historique de validation ici pour le correctif.

5
Louis T

Si tu veux train_test_split pour se comporter comme prévu (stratification par plusieurs colonnes sans doublons), créez une nouvelle colonne qui est une concaténation des valeurs dans vos autres colonnes et stratifiez sur la nouvelle colonne.

df['bc'] = df['b'].astype(str) + df['c'].astype(str)
train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['bc']])

Si vous craignez une collision en raison de valeurs telles que 11 et 3 et 1 et 13 les deux créant une valeur concaténée de 113, alors vous pouvez ajouter une chaîne arbitraire au milieu:

df['bc'] = df['b'].astype(str) + "_" + df['c'].astype(str)
3
Sesquipedalism