Dividindo o dataset em treino e teste no Python

Quem está começando os primeiros modelos estatísticos sabe que uma coisa padrão é dividir a base em duas: treino e teste. A primeira é utilizada para construção do modelo e a segunda para testá-lo, verificar sua performance e entender se ele está overfittando – para overfitting, ler Overfitting e Explicando overfitting com uma anedota. Hoje, vamos ver como fazer essa divisão no Python. Bora!

Divisão utilizando índices

Antes de mais nada, começaremos com a divisão através dos índices do dataset. Aqui, você deve sortear posições de forma aleatória utilizando o Numpy e, a partir delas, você seleciona os datasets. Veja um exemplo com a divisão 80-20:

import numpy as np

msk = np.random.rand(len(df)) < 0.8 # cria uma mascara com as posicoes
train = df[msk]
test = df[~msk]

Divisão utilizando função do sklearn

Adicionalmente, podemos utilizar a função train_test_split() do scikit-learn, a biblioteca mais popular em ciência de dados utilizando Python:

from sklearn.model_selection import train_test_split
train, test = train_test_split(df, train_size=0.8)

Não só podemos dividir entre treino e teste, como também podemos utilizar a divisão explicitando as variáveis resposta e explicativa. Só para exemplificar, poderíamos fazer o seguinte:

from sklearn.model_selection import train_test_split 
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8)

Sem dúvida, a função do sklearn é muito mais prática. Note que a sintaxe é bem simples, você passa o nome da sua tabela e a proporção do treino. Em resumo, você deve se atentar aos seguintes parâmetros:

• test_size: tamanho do teste, que no nosso caso seria 0.2;
• random_state: controla o embaralhamento dos dados. Se você utilizasse um inteiro aqui, essa seleção de treino e teste seria replicável;
• shuffle: recebe um booleano, True ou False, para determinar se os dados devem ser embaralhados antes da divisão;
• stratify: determina se os dados devem ser estratificados de acordo com alguma classe. Ex.: você quer uma divisão semelhante de homens e mulheres no treino e teste. Recebe como parâmetro essa classe.

Dessa maneira, conseguimos reduzir os riscos de overfitting no nosso modelo. Uma linha de código e nossos problemas estão (quase) solucionados!

Gostou do conteúdo? Se inscreva para receber as novidades! Deixe seu e-mail em INSCREVA-SE na barra à direita, logo abaixo de pesquisar. E, por favor, não deixe de comentar, dar seu feedback e compartilhar com seus amigos. De verdade, isso faz toda a diferença. Você também pode acompanhar mais do meu trabalho seguindo a conta de Twitter @EstatSite ou por alguma das redes que você encontra em Sobre o Estatsite / Contato.

Bons estudos!