0Pricing
Machine Learning Academy · Lesson

train_test_split: Ratios, Seeds, and Stratification

Learners will use sklearn's train_test_split with various test sizes, set random seeds for reproducibility, and apply stratification to balanced splits.

train_test_split: Ratios, Seeds, and Stratification is a free Machine Learning Academy lesson on CoddyKit — lesson 2 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Machine Learning Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

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.

Frequently asked questions

Is the “train_test_split: Ratios, Seeds, and Stratification” lesson free?

Yes — the full text of “train_test_split: Ratios, Seeds, and Stratification” is free to read here on the web, and the Machine Learning Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Machine Learning Academy course, upgrade to CoddyKit PRO.

What will I learn in “train_test_split: Ratios, Seeds, and Stratification”?

Learners will use sklearn's train_test_split with various test sizes, set random seeds for reproducibility, and apply stratification to balanced splits. You practise Machine Learning Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start Machine Learning Academy?

No prior experience is required. Machine Learning Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 2 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “train_test_split: Ratios, Seeds, and Stratification” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this Machine Learning Academy lesson?

Yes. Every Machine Learning Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

  1. Why You Cannot Evaluate on Training Data
  2. train_test_split: Ratios, Seeds, and Stratification
  3. Bias-Variance Trade-off: Underfitting vs Overfitting
  4. Baseline Models: Always Beat the Dummy Classifier
← Back to Machine Learning Academy