Treino estratificado / divisão de teste no scikit-learn

88

Preciso dividir meus dados em um conjunto de treinamento (75%) e um conjunto de teste (25%). Atualmente faço isso com o código abaixo:

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

No entanto, gostaria de estratificar meu conjunto de dados de treinamento. Como faço isso? Estive examinando o StratifiedKFoldmétodo, mas não me permite especificar a divisão de 75% / 25% e apenas estratificar o conjunto de dados de treinamento.

pir
fonte

Respostas:

153

[atualização para 0,17]

Veja a documentação 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)

[/ atualização para 0,17]

Há uma solicitação de pull aqui . Mas você pode simplesmente fazer train, test = next(iter(StratifiedKFold(...))) e usar os índices de treinar e testar se quiser.

Andreas Mueller
fonte
1
@AndreasMueller Existe uma maneira fácil de estratificar dados de regressão?
Jordan,
3
@Jordan nada é implementado no scikit-learn. Não conheço uma forma padrão. Podemos usar percentis.
Andreas Mueller
@AndreasMueller Você já viu o comportamento em que esse método é consideravelmente mais lento do que o StratifiedShuffleSplit? Eu estava usando o conjunto de dados MNIST.
snymkpr
@activatedgeek que parece muito estranho, já que train_test_split (... stratify =) está apenas chamando StratifiedShuffleSplit e fazendo a primeira divisão. Sinta-se à vontade para abrir um problema no rastreador com um exemplo reproduzível.
Andreas Mueller,
@AndreasMueller Na verdade, não abri um problema porque tenho a forte sensação de que estou fazendo algo errado (embora sejam apenas 2 linhas). Mas se ainda for capaz de reproduzi-lo várias vezes hoje, farei isso!
snymkpr
29

TL; DR: Use StratifiedShuffleSplit comtest_size=0.25

O Scikit-learn oferece dois módulos para divisão estratificada:

  1. StratifiedKFold : Este módulo é útil como um operador de validação cruzada k-fold direto: conforme nele, irá configurar conjuntos de n_foldstreinamento / teste de forma que as classes sejam igualmente balanceadas em ambos.

Aqui está algum código (diretamente da documentação acima)

>>> 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 : Este módulo cria um único conjunto de treinamento / teste com classes igualmente balanceadas (estratificadas). Essencialmente, é isso que você deseja com o n_iter=1. Você pode mencionar o tamanho do teste aqui mesmo que emtrain_test_split

Código:

>>> 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
picante
fonte
5
Observe que a partir de 0.18.x, n_iterdeve ser n_splitspara StratifiedShuffleSplit - e que há uma API ligeiramente diferente para isso: scikit-learn.org/stable/modules/generated/…
lollercoaster
2
Se yfor uma série Pandas, usey.iloc[train_index], y.iloc[test_index]
Owlright 01 de
1
@Owlright Eu tentei usar um dataframe do pandas e os índices que StratifiedShuffleSplit retorna não são os índices no dataframe. dataframe index: 2,3,5 the first split in sss:[(array([2, 1]), array([0]))]:(
Meghna Natraj
2
@tangy por que isso é um loop for? não é o caso de quando uma linha X_train, X_test = X[train_index], X[test_index]é chamada, ela se sobrepõe X_traine X_test? Por que então não apenas um next(sss)?
Bartek Wójcik
13

Aqui está um exemplo para dados contínuos / regressão (até que esse problema no GitHub seja resolvido).

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
)
  • Onde starté mínimo e stopmáximo de sua meta contínua.
  • Se você não definir right=True, isso fará com que seu valor máximo seja mais ou menos um compartimento separado e sua divisão sempre falhará porque poucas amostras estarão nesse compartimento extra.
Jordânia
fonte
12

Você pode simplesmente fazer isso com o train_test_split()método disponível no 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']) 

Também preparei um breve GitHub Gist que mostra como a stratifyopção funciona:

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

Shayan Amani
fonte
6

Além da resposta aceita por @Andreas Mueller, só quero adicionar como @tangy mencionado acima:

StratifiedShuffleSplit se assemelha mais a train_test_split ( stratify = y) com recursos adicionais de:

  1. estratificar por padrão
  2. especificando n_splits , ele divide os dados repetidamente
Max
fonte
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)
José carlos castro
fonte
0

Atualizando a resposta @tangy acima para a versão atual do scikit-learn: 0.23.2 ( documentação 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]
Roei Bahumi
fonte