Machine Learning: Treinamento e Teste em Python

A divisão treino/teste é o passo mais fundamental na construção de um modelo de machine learning. A ideia é simples: mantém parte dos seus dados oculta do modelo durante o treinamento e, em seguida, mede o desempenho do modelo nessa porção oculta. Sem essa separação, não há como saber honestamente se o modelo aprendeu genuinamente um padrão ou simplesmente memorizou os exemplos de treinamento.

Este capítulo aborda como funciona o train_test_split do scikit-learn, o que cada parâmetro faz e como avaliar o modelo resultante — tanto para problemas de regressão quanto de classificação.

Por Que Separar os Dados de Treinamento dos de Teste?

Um modelo que treina e avalia nos mesmos dados parecerá muito mais preciso do que realmente é. Isso é chamado de vazamento de dados ou overfitting: o modelo memorizou os exemplos de treinamento em vez de aprender um padrão generalizável.

Imagine um estudante que estuda 100 questões de prática e depois faz um teste usando essas mesmas 100 questões. Ele pode tirar 100% — mas essa nota não diz nada sobre se ele realmente entende o assunto.

A divisão treino/teste evita isso ao:

• Treinar o modelo em uma parte dos dados para que ele possa aprender padrões.
• Testar o modelo em uma parte separada, não vista, para medir o desempenho no mundo real.

Uma divisão típica é 80% para treino / 20% para teste, embora a proporção ideal dependa da quantidade de dados disponível.

A Função train_test_split

O scikit-learn fornece train_test_split no módulo model_selection. Ele embaralha aleatoriamente o conjunto de dados e o divide em duas partes:

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, random_state=42
)

Os quatro valores de retorno estão sempre nesta ordem: features de treino, features de teste, rótulos de treino, rótulos de teste.

Parâmetros Principais

Efeito de test_size

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris

X, y = load_iris(return_X_y=True)   # 150 samples

for ts in [0.1, 0.2, 0.3]:
    X_tr, X_te, _, _ = train_test_split(X, y, test_size=ts, random_state=42)
    print(f"test_size={ts}: train={len(X_tr)}, test={len(X_te)}")

Saída:

test_size=0.1: train=135, test=15
test_size=0.2: train=120, test=30
test_size=0.3: train=105, test=45

random_state e Reprodutibilidade

Sem um random_state, a divisão é diferente toda vez que o script é executado. Defina-o para qualquer inteiro para obter resultados reproduzíveis:

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(10).reshape(-1, 1)
y = np.arange(10)

_, X_te1, _, _ = train_test_split(X, y, test_size=0.3, random_state=42)
_, X_te2, _, _ = train_test_split(X, y, test_size=0.3, random_state=42)

print("Same random_state → same split:", list(X_te1.ravel()) == list(X_te2.ravel()))
# True

A escolha do inteiro (42, 0, 1, etc.) não importa — desde que você use o mesmo valor de forma consistente.

Exemplo de Regressão: Prevendo Preços de Imóveis

O exemplo a seguir gera um conjunto de dados sintético de tamanhos e preços de imóveis, treina um modelo de regressão linear e o avalia no conjunto de teste retido.

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

# Generate synthetic data: house size (sq ft) → price
np.random.seed(42)
n = 200
sqft = np.random.randint(500, 3500, n).astype(float)
price = 150 * sqft + np.random.randn(n) * 20000

X = sqft.reshape(-1, 1)
y = price

# Split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print("Training samples:", len(X_train))   # 160
print("Testing samples:", len(X_test))     # 40

# Train
model = LinearRegression()
model.fit(X_train, y_train)

# Evaluate on the test set
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
r2  = r2_score(y_test, y_pred)

print(f"Mean Squared Error: {mse:,.0f}")
print(f"R² Score:           {r2:.4f}")
print(f"Coefficient:        {model.coef_[0]:.2f}")
print(f"Intercept:          {model.intercept_:.2f}")

Saída:

Training samples: 160
Testing samples: 40
Mean Squared Error: 489,271,263
R² Score:           0.9651
Coefficient:        147.55
Intercept:          5237.02

Interpretando as métricas:

• MSE (Erro Quadrático Médio) é a diferença quadrática média entre previsões e valores reais. Quanto menor, melhor, mas a escala depende da variável alvo (aqui, dólares).
• R² varia de 0 a 1. Um valor de 0,965 significa que o modelo explica cerca de 96,5% da variância nos preços dos imóveis — um bom ajuste para esse conjunto de dados simples.

Para mais informações sobre regressão linear, veja Regressão Linear em Python.

Exemplo de Classificação: Espécies de Flores Iris

Para uma tarefa de classificação, a acurácia e um relatório de classificação são mais informativos do que o MSE.

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report

X, y = load_iris(return_X_y=True)

# stratify=y ensures each class is proportionally represented in both splits
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print("Train class distribution:", np.bincount(y_train))  # [40 40 40]
print("Test class distribution: ", np.bincount(y_test))   # [10 10 10]

model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)
print("\nAccuracy:", round(accuracy_score(y_test, y_pred), 4))
print()
print(classification_report(y_test, y_pred,
      target_names=["setosa", "versicolor", "virginica"]))

Saída:

Train class distribution: [40 40 40]
Test class distribution:  [10 10 10]

Accuracy: 0.9667

              precision    recall  f1-score   support

      setosa       1.00      1.00      1.00        10
  versicolor       1.00      0.90      0.95        10
   virginica       0.91      1.00      0.95        10

    accuracy                           0.97        30
   macro avg       0.97      0.97      0.97        30
weighted avg       0.97      0.97      0.97        30

Por que stratify=y importa: Sem ele, uma divisão aleatória de um conjunto de dados desbalanceado poderia colocar a maioria das amostras de uma classe rara no treinamento, não deixando nenhuma no conjunto de teste. stratify=y garante que cada classe apareça nas mesmas proporções em ambas as divisões.

Para mais informações sobre classificação, veja Regressão Logística em Python e Matriz de Confusão em Python.

Armadilhas Comuns

Pré-processamento Antes ou Depois da Divisão?

Sempre divida os dados antes de ajustar qualquer etapa de pré-processamento (escalonamento, codificação, imputação). Se você escalonar todos os dados e depois dividir, o conjunto de teste terá influenciado o escalonador — uma forma de vazamento de dados.

A ordem correta:

from sklearn.preprocessing import StandardScaler
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, random_state=42)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)  # fit ONLY on training data
X_test_scaled  = scaler.transform(X_test)       # transform test with same parameters

Veja Escalonamento de Features em Python para um guia completo.

Embaralhamento e Dados de Séries Temporais

train_test_split embaralha os dados por padrão. Para dados de séries temporais isso está errado — você estaria treinando em dados futuros para prever o passado. Defina shuffle=False e certifique-se de que seus dados estejam ordenados cronologicamente antes de dividir:

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, shuffle=False
)

Quando Uma Única Divisão Não É Suficiente

Uma única divisão 80/20 fornece uma estimativa de desempenho do modelo que depende de quais amostras foram parar no conjunto de teste. A validação cruzada repete a divisão várias vezes e calcula a média das pontuações, fornecendo uma estimativa muito mais estável — especialmente em conjuntos de dados pequenos.

Escolhendo uma Métrica de Avaliação

A métrica correta depende do tipo de problema:

Para classificação binária, veja Curva AUC-ROC em Python e Matriz de Confusão em Python. Para ajuste de hiperparâmetros após ter uma divisão funcionando, veja Grid Search em Python.

Resumo

• Divida seus dados antes de qualquer pré-processamento para evitar vazamento de dados.
• Use test_size=0.2 como padrão razoável; ajuste conforme o tamanho do conjunto de dados.
• Defina random_state para qualquer inteiro para obter divisões reproduzíveis.
• Use stratify=y para tarefas de classificação, especialmente em dados desbalanceados.
• Defina shuffle=False para dados de séries temporais.
• Uma única divisão treino/teste é uma linha de base rápida; use validação cruzada para estimativas de desempenho mais confiáveis.