Skip to content

Sklearn train_test_split: Dividir Dados para Machine Learning em Python

Soren Atelier
Name
Soren Atelier

Updated on

Treinar um modelo de machine learning com os mesmos dados usados para avaliá-lo produz uma acurácia enganosamente alta. O modelo memoriza os dados de treino em vez de aprender padrões generalizáveis -- um problema chamado overfitting (sobreajuste). Você precisa de um conjunto de teste separado que o modelo nunca veja durante o treino para obter métricas de desempenho honestas.

A função train_test_split() do Scikit-learn é a forma padrão de dividir conjuntos de dados em porções de treino e teste. Ela lida com arrays, DataFrames e matrizes esparsas, com opções para estratificação, reprodutibilidade e proporções de divisão personalizadas.

📚

    Uso básico

    from sklearn.model_selection import train_test_split
import numpy as np
 
# Dados de exemplo: 100 amostras, 5 características
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)  # Labels binários
 
# Divisão: 80% treino, 20% teste
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
 
print(f"Conjunto de treino: {X_train.shape[0]} amostras")
print(f"Conjunto de teste:  {X_test.shape[0]} amostras")
# Conjunto de treino: 80 amostras
# Conjunto de teste:  20 amostras

    Parâmetros principais

    ParâmetroPadrãoDescrição
    test_size0.25Fração (0.0-1.0) ou número absoluto de amostras de teste
    train_sizeNoneFração ou número de amostras de treino (complemento de test_size)
    random_stateNoneSemente para divisões reprodutíveis
    shuffleTrueSe os dados devem ser embaralhados antes da divisão
    stratifyNoneArray para divisão estratificada

    Com Pandas DataFrames

    from sklearn.model_selection import train_test_split
import pandas as pd
 
df = pd.DataFrame({
    'age': [25, 30, 35, 40, 45, 50, 55, 60, 28, 33],
    'income': [40, 50, 60, 70, 80, 90, 100, 110, 45, 55],
    'purchased': [0, 0, 1, 1, 1, 1, 1, 1, 0, 0]
})
 
X = df[['age', 'income']]
y = df['purchased']
 
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)
 
print(f"Train shape: {X_train.shape}")  # (7, 2)
print(f"Test shape:  {X_test.shape}")   # (3, 2)
print(type(X_train))  # <class 'pandas.core.frame.DataFrame'>

    DataFrames permanecem como DataFrames após a divisão -- nomes de colunas e índices são preservados.

    random_state: Divisões reprodutíveis

    Sem random_state, você obtém uma divisão diferente a cada execução:

    from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.arange(10).reshape(5, 2)
y = np.array([0, 0, 1, 1, 1])
 
# Sem random_state: divisão diferente a cada execução
_, X_test1, _, _ = train_test_split(X, y, test_size=0.4)
_, X_test2, _, _ = train_test_split(X, y, test_size=0.4)
print(np.array_equal(X_test1, X_test2))  # Provavelmente False
 
# Com random_state: mesma divisão toda vez
_, X_test3, _, _ = train_test_split(X, y, test_size=0.4, random_state=42)
_, X_test4, _, _ = train_test_split(X, y, test_size=0.4, random_state=42)
print(np.array_equal(X_test3, X_test4))  # True

    Sempre defina random_state para reprodutibilidade. Use qualquer inteiro -- 42 é convencional, mas o número específico não importa.

    Divisão estratificada

    Para conjuntos de dados desbalanceados, uma divisão aleatória pode colocar a maioria das amostras minoritárias em um único conjunto. A estratificação garante que ambos os conjuntos tenham as mesmas proporções de classe:

    from sklearn.model_selection import train_test_split
import numpy as np
from collections import Counter
 
# Conjunto de dados desbalanceado: 90% classe 0, 10% classe 1
np.random.seed(42)
X = np.random.randn(200, 4)
y = np.array([0] * 180 + [1] * 20)
 
# Sem estratificação
_, _, y_train_bad, y_test_bad = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print("Sem estratificação:")
print(f"  Train: {Counter(y_train_bad)}")
print(f"  Test:  {Counter(y_test_bad)}")
 
# Com estratificação
_, _, y_train_good, y_test_good = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print("\nCom stratify=y:")
print(f"  Train: {Counter(y_train_good)}")
print(f"  Test:  {Counter(y_test_good)}")
# Train: Counter({0: 144, 1: 16})  -- 10% classe 1
# Test:  Counter({0: 36, 1: 4})    -- 10% classe 1

    Quando usar estratificação

    CenárioUsar stratify?
    Classes balanceadas (50/50)Opcional
    Classes desbalanceadas (90/10)Sim
    Classificação multiclasseSim
    Regressão (alvo contínuo)Não (não suportado)
    Conjuntos pequenos (< 100 amostras)Sim (previne classes vazias)

    Divisão Treino/Validação/Teste

    Para ajuste de hiperparâmetros, você precisa de três conjuntos: treino, validação e teste. Aplique train_test_split duas vezes:

    from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.random.randn(1000, 10)
y = np.random.randint(0, 3, 1000)
 
# Primeira divisão: 80% treino+val, 20% teste
X_temp, X_test, y_temp, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# Segunda divisão: 75% treino, 25% val (dos 80% = 60/20/20 no total)
X_train, X_val, y_train, y_val = train_test_split(
    X_temp, y_temp, test_size=0.25, random_state=42, stratify=y_temp
)
 
print(f"Treino:     {X_train.shape[0]} amostras (60%)")
print(f"Validação:  {X_val.shape[0]} amostras (20%)")
print(f"Teste:      {X_test.shape[0]} amostras (20%)")

    Exemplo completo de pipeline ML

    from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
import numpy as np
 
# Gerar dados de exemplo
np.random.seed(42)
X = np.random.randn(500, 8)
y = (X[:, 0] + X[:, 1] > 0).astype(int)
 
# Dividir dados
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# Escalar características (fit apenas nos dados de treino!)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)  # Usar estatísticas do treino
 
# Treinar modelo
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
 
# Avaliar
y_pred = model.predict(X_test_scaled)
print(f"Acurácia: {accuracy_score(y_test, y_pred):.4f}")
print(classification_report(y_test, y_pred))

    Regra crítica: Ajuste o scaler (e qualquer pré-processamento) apenas no conjunto de treino. Aplique a mesma transformação ao conjunto de teste. Ajustar no conjunto completo causa vazamento de dados (data leakage).

    Proporções de divisão comuns

    DivisãoTreinoTesteQuando usar
    80/2080%20%Escolha padrão, maioria dos conjuntos
    70/3070%30%Conjuntos pequenos, necessidade de conjunto de teste maior
    90/1090%10%Conjuntos grandes (10k+ amostras)
    60/20/2060%20% val + 20% testeAo ajustar hiperparâmetros

    Explorando resultados do modelo

    Após dividir e treinar seu modelo, o PyGWalker (opens in a new tab) permite explorar interativamente previsões vs valores reais, distribuições de características entre conjuntos de treino/teste e padrões de erro no Jupyter:

    import pandas as pd
import pygwalker as pyg
 
results = pd.DataFrame({
    'actual': y_test,
    'predicted': y_pred,
    'correct': y_test == y_pred
})
walker = pyg.walk(results)

    FAQ

    O que o train_test_split faz no sklearn?

    train_test_split() divide aleatoriamente arrays ou DataFrames em dois subconjuntos: um para treino e um para teste. Garante que a avaliação do modelo use dados que ele não viu durante o treino, fornecendo estimativas de desempenho honestas.

    Qual é a melhor proporção de divisão treino/teste?

    80/20 é o padrão. Use 70/30 para conjuntos menores onde você precisa de um conjunto de teste confiável, e 90/10 para conjuntos grandes (10k+ amostras) onde 10% ainda é substancial. Para ajuste de hiperparâmetros, use 60/20/20 (treino/val/teste).

    O que o random_state faz no train_test_split?

    random_state define a semente aleatória para o embaralhamento que ocorre antes da divisão. Usar o mesmo valor de random_state produz a mesma divisão toda vez, tornando seus resultados reprodutíveis. Qualquer inteiro funciona.

    Quando devo usar stratify no train_test_split?

    Use stratify=y quando sua variável alvo é desbalanceada (ex., 95% negativo, 5% positivo) ou quando você tem um conjunto de dados pequeno. A estratificação garante que os conjuntos de treino e teste tenham a mesma proporção de cada classe.

    Como divido dados em conjuntos de treino, validação e teste?

    Chame train_test_split duas vezes. Primeiro divida em treino+val e teste (ex., 80/20). Depois divida treino+val em treino e val (ex., 75/25 dos 80%, resultando em 60/20/20 no total). Alternativamente, use sklearn.model_selection.KFold para validação cruzada.

    Conclusão

    train_test_split() é a base de todo fluxo de trabalho de machine learning em Python. Sempre use antes de treinar -- nunca avalie com dados de treino. Defina random_state para reprodutibilidade, use stratify=y para classes desbalanceadas, e lembre-se de ajustar etapas de pré-processamento apenas no conjunto de treino. Para seleção de modelos, divida em três partes (treino/validação/teste) ou use validação cruzada para estimativas mais robustas.

    📚