Como dividir dados em treino e teste em Machine Learning

Dividir os dados em treino e teste é um dos primeiros passos de qualquer projeto de Machine Learning: treina-se o modelo num conjunto e avalia-se noutro que ele nunca viu, evitando resultados demasiado otimistas. A seguir vai ver como dividir dados em treino e teste em Machine Learning com a função train_test_split do scikit-learn, com exemplos simples que pode adaptar ao seu caso.

Pré-requisitos

• Python 3 instalado (versão 3.9 ou superior).
• As bibliotecas scikit-learn e pandas instaladas (pip install scikit-learn pandas).
• Um conjunto de dados com colunas de entrada (features) e uma coluna-alvo (target).
• Noções básicas de Python, como variáveis e importação de módulos.

Passo 1: Preparar os dados

Antes de dividir, separamos as features (as colunas que o modelo usa para prever, guardadas em X) da coluna-alvo (aquilo que queremos prever, guardado em y). Vamos usar um exemplo pequeno com pandas.

import pandas as pd

# Example dataset
data = pd.DataFrame({
    "study_hours": [1, 2, 3, 4, 5, 6, 7, 8],
    "absences":    [8, 6, 5, 4, 3, 2, 1, 0],
    "passed":      [0, 0, 0, 1, 1, 1, 1, 1]
})

X = data[["study_hours", "absences"]]  # features
y = data["passed"]                      # target

Por convenção, X é uma tabela com várias colunas e y é uma única coluna. Manter esta separação clara evita que a coluna-alvo entre por engano nas features — um erro que faria o modelo "espreitar" a resposta.

Passo 2: Dividir com train_test_split

A função train_test_split do módulo sklearn.model_selection baralha os dados e separa-os em dois conjuntos. Basta passar X e y e indicar a percentagem para teste.

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,     # 20% for testing, 80% for training
    random_state=42    # makes the split reproducible
)

Repare na ordem das variáveis à esquerda: primeiro os dois conjuntos de X (treino e teste) e só depois os dois de y. Trocar esta ordem é um erro comum que provoca resultados estranhos mais à frente.

Passo 3: Entender test_size e random_state

Dois parâmetros controlam o comportamento da divisão:

• test_size: a fração dos dados reservada para teste. 0.2 significa 20%; em alternativa pode indicar train_size.
• random_state: fixa a semente aleatória. Usando sempre o mesmo número (por exemplo 42), obtém exatamente a mesma divisão de cada vez que corre o código, tornando as experiências reproduzíveis.

Por defeito, a função baralha as linhas antes de dividir (shuffle=True). Isto é importante porque muitos conjuntos de dados vêm ordenados — por exemplo, com todos os casos positivos no fim — e uma divisão sem baralhar deixaria o teste enviesado.

Passo 4: Manter as proporções com stratify

Em problemas de classificação, treino e teste devem ter a mesma proporção de classes. Se 30% dos casos são da classe "passou", queremos esses 30% em ambos os conjuntos. O parâmetro stratify garante isso:

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42,
    stratify=y         # keep class proportions of y
)

Dica: use stratify=y sempre que a variável-alvo for categórica. Em problemas de regressão (alvo contínuo), deixe stratify de fora.

Verificar o resultado

Para confirmar que a divisão correu bem, veja o número de linhas de cada conjunto com o atributo .shape e as proporções de classes:

print("Train:", X_train.shape, "Test:", X_test.shape)
print(y_train.value_counts(normalize=True))

A soma das linhas de treino e teste deve ser igual ao total original, e as proporções de classes devem ficar semelhantes nos dois conjuntos. Como fixámos o random_state, estes números repetem-se sempre que corre o script.

Conclusão

Com poucas linhas separou os dados em treino e teste — o passo que garante uma avaliação honesta de um modelo de Machine Learning. A partir daqui pode treinar um algoritmo com X_train e y_train e medir o desempenho com X_test e y_test, por exemplo com uma matriz de confusão. Como próximo passo, experimente mudar o test_size para 0.3 e observe como muda o tamanho dos conjuntos. Qual será a percentagem de teste ideal para o seu problema?