0Pricing
Machine Learning Academy · Aula

train_test_split: proporções, sementes e estratificação

Use train_test_split do sklearn com vários tamanhos de teste, defina sementes aleatórias para garantir a reprodutibilidade e aplique estratificação para obter divisões equilibradas.

train_test_split: proporções, sementes e estratificação é uma aula grátis de Machine Learning Academy no CoddyKit. Esta é a aula 2 de 4. Você pode ler a aula completa abaixo gratuitamente — depois pratica ao vivo no navegador com um editor de código integrado e um tutor de IA 24/7. Faz parte do caminho de aprendizado de Machine Learning Academy, e seu progresso é sincronizado entre a web e o app CoddyKit. O curso de Machine Learning Academy inclui 4 aulas no total.

Partes desta aula ainda não foram traduzidas e aparecem em inglês.

The train_test_split Function

Scikit-learn's train_test_split() function is the standard tool for partitioning datasets into training and test subsets. It randomly shuffles the data and allocates the specified proportion to the test set, returning four arrays: X_train, X_test, y_train, y_test.

Understanding its parameters fully is important because poor splitting choices can compromise the validity of all subsequent evaluation. A too-small test set gives noisy estimates; a too-large test set wastes training data. The right choice depends on your dataset size and the variance you can afford in performance estimates.

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

X, y = load_breast_cancer(return_X_y=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.20,     # 20% of data goes to test set
    random_state=42,    # reproducible random shuffle
    shuffle=True,       # default: shuffle before splitting
    stratify=y          # preserve class proportions
)

print(f'Training size: {X_train.shape}')
print(f'Test size:     {X_test.shape}')

Choosing the Test Size Ratio

The test_size parameter accepts either a float (proportion) or an integer (absolute count):

  • Large datasets (>100k examples): use 10-15% test (there is plenty of training data; a large test set gives very stable estimates).
  • Medium datasets (1k–100k): 20-30% test is standard (80/20 or 70/30 splits).
  • Small datasets (<1k): consider cross-validation instead of a single split, because any fixed split may be unlucky.

There is no universally correct ratio. The goal is enough test examples for a statistically stable performance estimate, while keeping enough training data for the model to learn well.

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_breast_cancer
from sklearn.linear_model import LogisticRegression
import numpy as np

X, y = load_breast_cancer(return_X_y=True)

# Compare different test sizes
for test_size in [0.1, 0.2, 0.3, 0.4]:
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=42)
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)
    acc = model.score(X_test, y_test)
    print(f'test_size={test_size}: train={len(X_train)}, test={len(X_test)}, acc={acc:.3f}')

Random Seeds for Reproducibility

The random_state parameter seeds the random number generator that shuffles the data before splitting. With the same seed, the split is identical every time. Without a seed, the split changes on every run.

Reproducibility matters for: comparing models fairly (same split for all), sharing code with colleagues, debugging (same data each run), and scientific publication (results can be verified). Use a fixed seed for all experiments, but try two or three different seeds to confirm your model is not accidentally benefiting or suffering from one lucky split.

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer

X, y = load_breast_cancer(return_X_y=True)

# Test variance across random seeds
results = []
for seed in [0, 1, 42, 123, 999]:
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=seed
    )
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)
    acc = model.score(X_test, y_test)
    results.append(acc)
    print(f'seed={seed}: accuracy={acc:.4f}')

import numpy as np
print(f'\nMean: {np.mean(results):.4f}  Std: {np.std(results):.4f}')

The Problem Without Stratification

Without stratification, random splitting may produce a test set where the class proportions differ significantly from the overall dataset. If your dataset has 90% class 0 and 10% class 1, a random split might give a test set with only 3% class 1 by chance.

This matters because: performance metrics computed on an atypical class distribution are not representative of production performance. A model might score perfectly on a test set that contains almost no minority-class examples, giving false confidence in a model that actually fails on the positive class.

from sklearn.model_selection import train_test_split
import numpy as np

# Imbalanced dataset: 90% class 0, 10% class 1
np.random.seed(0)
n = 1000
y_imb = np.array([0]*900 + [1]*100)

results = []
for seed in range(10):
    _, _, y_train, y_test = train_test_split(y_imb, y_imb, test_size=0.2, random_state=seed)
    results.append(y_test.mean())
    print(f'seed={seed}: test positive rate = {y_test.mean():.3f} (expected: 0.100)')

print(f'\nMin: {min(results):.3f}  Max: {max(results):.3f}')

Stratification: Preserving Class Proportions

Setting stratify=y ensures that both the training and test sets contain the same proportion of each class as the full dataset. This is essential for imbalanced classification problems and should be used by default whenever you have a classification target.

Stratification works by splitting each class separately and then combining: if your dataset has 90% class 0 and 10% class 1, both the training set and test set will also have approximately 90%/10% proportions. This makes performance estimates on the test set much more stable and representative.

from sklearn.model_selection import train_test_split
import numpy as np

np.random.seed(0)
n = 1000
y_imb = np.array([0]*900 + [1]*100)
X_imb = np.random.randn(n, 5)

# Without stratification
_, _, y_train_no, y_test_no = train_test_split(X_imb, y_imb, test_size=0.2, random_state=42)
print(f'Without stratify: test positive rate = {y_test_no.mean():.3f}')

# With stratification
_, _, y_train_s, y_test_s = train_test_split(X_imb, y_imb, test_size=0.2, random_state=42, stratify=y_imb)
print(f'With stratify:    test positive rate = {y_test_s.mean():.3f}  (exactly 0.100)')

Stratification in Multi-Class Problems

Stratification works equally well for multi-class targets. When you pass stratify=y with a multi-class y, scikit-learn ensures all classes are represented in the same proportions in both splits.

This is particularly important for rare classes. If class C appears in only 2% of examples and you have a 500-example dataset, a random split might put all 10 class-C examples in training and none in the test set — making the test set useless for evaluating performance on class C. Stratification prevents this.

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

X, y = load_iris(return_X_y=True)  # 3 balanced classes

# Stratified split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)

# Verify class proportions are preserved
for cls in [0, 1, 2]:
    full_ratio = (y == cls).mean()
    train_ratio = (y_train == cls).mean()
    test_ratio = (y_test == cls).mean()
    print(f'Class {cls}: full={full_ratio:.2f} train={train_ratio:.2f} test={test_ratio:.2f}')

Shuffling: Why Order Matters

By default, shuffle=True randomly reorders the data before splitting. This is critical because real datasets are often not in random order — they may be sorted by date, user ID, or label. Without shuffling, the test set would contain only the last examples in the dataset, which may be systematically different from earlier examples.

The exception is time-series data, where you should NOT shuffle. For time-ordered data, the test set must come from the future relative to training data. Setting shuffle=False preserves time order. Scikit-learn's TimeSeriesSplit provides more sophisticated time-ordered cross-validation.

from sklearn.model_selection import train_test_split
import numpy as np

# Time-series data: ordered by date
dates = np.arange(1000)           # day 0 to 999
sales = dates * 0.5 + np.random.randn(1000) * 10

# WRONG for time series: shuffling leaks future data into training
X_train_wrong, X_test_wrong, _, _ = train_test_split(dates.reshape(-1,1), sales, shuffle=True)
print(f'Wrong: test date range = {X_test_wrong.min():.0f} to {X_test_wrong.max():.0f}')

# CORRECT for time series: last 20% of time is the test set
X_train_right, X_test_right, _, _ = train_test_split(dates.reshape(-1,1), sales, shuffle=False)
print(f'Right: test date range = {X_test_right.min():.0f} to {X_test_right.max():.0f}')

Splitting Multiple Arrays Consistently

A common need is to split multiple arrays together — for example, X features, y labels, and a separate ID array — while keeping the same indices aligned. train_test_split accepts any number of arrays as positional arguments and applies the same shuffle and split to all of them.

This is also how you split pre-processed feature matrices alongside original DataFrames if you want to track which original rows ended up in which split for error analysis.

from sklearn.model_selection import train_test_split
import numpy as np

# Multiple arrays with same sample dimension
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)
sample_ids = np.arange(100)  # e.g., customer IDs

# Split all three together
X_train, X_test, y_train, y_test, ids_train, ids_test = train_test_split(
    X, y, sample_ids,
    test_size=0.2,
    random_state=42
)

print('Train IDs sample:', ids_train[:5])
print('Test IDs sample:', ids_test[:5])
print('All arrays remain aligned.')

When to Use Cross-Validation Instead

A single train-test split has high variance — the exact split can produce performance estimates that vary by several percentage points. Cross-validation averages performance across multiple splits, giving a more stable and reliable estimate.

Use cross-validation (not a single split) when:

  • Your dataset is small (fewer than a few thousand examples).
  • You want a reliable estimate with low variance.
  • You are comparing models or tuning hyperparameters.

Use a single split when you have very large datasets where cross-validation is computationally prohibitive.

from sklearn.model_selection import cross_val_score
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer
import numpy as np

X, y = load_breast_cancer(return_X_y=True)

# 5-fold cross-validation (more reliable than single split)
scores = cross_val_score(
    LogisticRegression(max_iter=1000),
    X, y,
    cv=5,
    scoring='accuracy'
)

print('CV scores:', scores.round(4))
print(f'Mean: {scores.mean():.4f}  Std: {scores.std():.4f}')
# Much more stable estimate than any single split

Complete Evaluation Workflow

The complete evaluation workflow that avoids all the pitfalls discussed in this lesson:

  1. Split off 20% as a final test set (stratified, random seed set).
  2. Use the remaining 80% for all development (training, cross-validation, hyperparameter tuning).
  3. Once development is complete and you have selected your final model, fit it on the entire 80% training data.
  4. Evaluate on the 20% test set exactly once and report that score.

Following this workflow gives you a trustworthy estimate of how your model will perform on future unseen data.

from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

X, y = load_breast_cancer(return_X_y=True)

# Step 1: Reserve test set FIRST
X_dev, X_test, y_dev, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# Step 2: Cross-validate on development set
pipeline = Pipeline([('scaler', StandardScaler()), ('clf', LogisticRegression(max_iter=1000))])
cv_scores = cross_val_score(pipeline, X_dev, y_dev, cv=5, scoring='accuracy')
print(f'Dev CV accuracy: {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})')

# Step 3: Fit final model on all dev data
pipeline.fit(X_dev, y_dev)

# Step 4: One-time final evaluation
print(f'Final test accuracy: {pipeline.score(X_test, y_test):.4f}')

Quick Check

Test your understanding of Machine Learning with Python concepts from this lesson.

Lesson Recap

In this lesson you learned: test_size controls the proportion of data held out for final evaluation, random_state seeds the shuffle for reproducibility and fair model comparison, and stratify=y preserves class proportions in both splits, which is critical for imbalanced datasets. Next up we study the bias-variance trade-off and learn to diagnose underfitting versus overfitting by plotting training and validation error curves.

Perguntas Frequentes

A aula “train_test_split: proporções, sementes e estratificação” é grátis?

Sim — o texto completo de “train_test_split: proporções, sementes e estratificação” é grátis para ler aqui na web. Para praticá-la interativamente (um editor de código integrado e um tutor de IA 24/7) e desbloquear o restante do curso de Machine Learning Academy, atualize para CoddyKit PRO. O curso de Machine Learning Academy inclui 4 aulas no total.

O que vou aprender em “train_test_split: proporções, sementes e estratificação”?

Use train_test_split do sklearn com vários tamanhos de teste, defina sementes aleatórias para garantir a reprodutibilidade e aplique estratificação para obter divisões equilibradas. Você pratica Machine Learning Academy com código prático que executa diretamente no navegador, e um tutor de IA 24/7 responde suas dúvidas enquanto trabalha na aula.

Preciso ter experiência prévia para começar Machine Learning Academy?

Nenhuma experiência prévia é necessária. Machine Learning Academy no CoddyKit é estruturado para alunos iniciantes até avançados, então você pode começar aqui ou desde o início e aprender no seu ritmo. Esta é a aula 2 de 4.

Quanto tempo leva a aula “train_test_split: proporções, sementes e estratificação”?

A maioria das aulas CoddyKit leva cerca de 5–10 minutos. Cada uma é compacta e interativa, então você faz progresso constante e retoma exatamente de onde parou entre web e app.

Posso escrever e executar código nesta aula de Machine Learning Academy?

Sim. Cada aula de Machine Learning Academy inclui um editor de código integrado, então você escreve e executa código real direto no navegador e recebe feedback de IA instantaneamente — nenhuma configuração local necessária.

Todas as aulas deste curso

  1. Por que não se pode avaliar com dados de treinamento
  2. train_test_split: proporções, sementes e estratificação
  3. Compromisso entre viés e variância: subajuste vs. sobreajuste
  4. Modelos de referência: supere sempre o DummyClassifier
← Voltar para Machine Learning Academy