Split Data into Train/Test Sets

You've just built a model that aces every question you threw at it — so you ship it, only to watch it stumble on real-world data. The culprit isn't your algorithm; it's that you evaluated it on the same data you trained it on. If you don't split your data into separate training and testing sets, your model is basically memorizing answers instead of learning patterns. This lesson fixes that: you'll master the train/test split — the single most important step before you ever fit a model — and set yourself up for honest, reliable model evaluation.

The problem this lesson solves

Imagine you're studying for a final exam by reading the exact same questions and answers over and over. On exam day, you'd ace it — but not because you understand the material. You just memorized the answers. That's precisely what happens when you train and evaluate your machine learning model on the same dataset. The model sees the same examples during training and testing, so it appears to perform brilliantly, but it fails completely on new, unseen data.

In data science, this is called overfitting — the model captures noise and quirks of the training set instead of the underlying pattern. The result is a model with near-perfect accuracy in your notebook but poor performance in production. The solution is to split your data into two disjoint subsets: a training set to teach the model, and a test set to evaluate how well it generalizes to data it has never seen.

Without a proper split, you cannot trust any metric (accuracy, precision, recall, RMSE) your model reports. This is a core tenet of the scientific method applied to machine learning: never evaluate a model on the data it was trained on. The train/test split is your first line of defense against self-deception, and it's a prerequisite for every supervised learning task, from regression to classification to deep learning.

Core concept / mental model

Think of a chef learning a new recipe. The chef practices with a training set of ingredients, adjusting seasoning and technique. Then, to test their skill, they cook the same dish using a fresh set of ingredients — the test set — and taste it. If the dish is great, you know the chef has truly learned the recipe, not just memorized one batch.

In technical terms, your dataset is a finite sample from a larger population. You want to build a model that performs well on that population, not just on the sample. By splitting your data, you simulate that larger population: the training set teaches the model, and the test set acts as a stand-in for unseen data. A model that performs well on both has generalized — it learned patterns that transfer, not noise that's specific to the training sample.

Python's scikit-learn library provides a simple, battle-tested function called train_test_split that does exactly this. It takes your features (X) and target (y), randomly shuffles them, and returns four arrays: X_train, X_test, y_train, y_test. Think of it as a robot chef that separates ingredients into practice portions and tasting portions.

Pro tip: The shuffle is crucial. If your data is ordered by time or class, a naive split (e.g., first 80% for training, last 20% for testing) can introduce bias. train_test_split shuffles by default, but you can control the randomness with the random_state parameter for reproducible results.

How it works step by step

Here is the logical flow of a train/test split, from raw data to model evaluation:

1. Load your dataset — usually as a pandas DataFrame or NumPy array.
2. Separate features (X) from target (y) — X contains all the predictor columns, y contains the column you want to predict.
3. Call train_test_split(X, y, test_size=0.2, random_state=42) — this shuffles the data and allocates 20% for testing and 80% for training.
4. Use X_train and y_train to fit your model — this is where the model learns its parameters.
5. Use X_test and y_test to evaluate — predict on X_test, compare to y_test, and compute metrics.

The test_size parameter controls the proportion of data held out for testing. Typical values range from 0.2 (20%) to 0.3 (30%). For very large datasets, you might use 0.1 (10% test). The random_state ensures the split is reproducible — run the same code again and get the exact same rows in each set.

Now let's see this in action with code.

Hands-on walkthrough

Setting up

First, make sure you have scikit-learn installed:

pip install scikit-learn

Basic train/test split

Let's create a small dataset and split it:

from sklearn.model_selection import train_test_split
import numpy as np

# Sample data: 100 samples, 2 features
X = np.random.rand(100, 2)
y = np.random.randint(0, 2, size=100)  # binary target

# Split into train and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

print(f"Full dataset: {X.shape} samples, {X.shape[1]} features")
print(f"Training set: {X_train.shape} samples")
print(f"Test set: {X_test.shape} samples")

Expected output:

Full dataset: (100, 2) samples, 2 features
Training set: (80, 2) samples
Test set: (20, 2) samples

Notice that the split is stratified by default if you pass y — it preserves the class proportions. For classification, that's important to ensure both sets have similar class ratios as the original dataset.

Using it with a real model

Here's a complete workflow with a logistic regression model:

from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

# Load a classic dataset
iris = load_iris()
X, y = iris.data, iris.target

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

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

# Evaluate on test set
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"Test accuracy: {accuracy:.2f}")

Expected output:

Test accuracy: 0.98

This accuracy is a honest estimate of how the model will perform on new iris flowers.

Splitting for regression

For regression, you don't need stratification (target is continuous), but you still split the same way:

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

# Load housing data
housing = fetch_california_housing()
X, y = housing.data, housing.target

# Split — no stratify here because y is continuous
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=7
)

# Train and evaluate
model = LinearRegression()
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f"Test MSE: {mse:.2f}")

Expected output:

Test MSE: 0.53

The MSE is calculated only on data the model never saw during training — the fair verdict.

Compare options / when to choose what

train_test_split is the go-to method for most projects, but it's not the only game in town. Here's when to use what:

For your first models, start with train_test_split. Once you want to tune hyperparameters, switch to cross-validation — but always keep a final held-out test set for the absolute last evaluation.

Stratification matters: for classification, always pass stratify=y to train_test_split to keep class proportions consistent. For regression, ignore it — it doesn't apply.

Troubleshooting & edge cases

Error: ValueError: n_samples mismatch

You get n_samples errors when X and y have different numbers of rows. Check that you sliced the DataFrame correctly:

# Wrong
X = df[['feature1', 'feature2']]
y = df['target']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
# Correct — make sure X has 2D shape and y is 1D

Always verify shapes before splitting:

print(X.shape, y.shape)

Mistake: Test set contaminated by train set

If you accidentally call train_test_split twice or mix the arrays, your test set may include training samples. Always use the four returned arrays separately — never reuse the original X or y for evaluation.

Overfitting despite split

If your model performs much worse on the test set than on the training set, that's a clear sign of overfitting. The train/test split surfaced it — that's a win! Next steps: simplify the model, add regularization, or gather more data.

random_state inconsistency

Without random_state, you'll get a different split each run, making debugging and reproducibility hard. Always set it, e.g., random_state=42, to get identical splits across runs and share with colleagues.

Imbalanced classification

If one class is rare (e.g., 1% of data), a random split might put all the rare class samples in the test set. Use stratify=y or use StratifiedShuffleSplit for more control.

What you learned & what's next

You can now explain the core idea behind splitting data into train and test sets: training on one subset and evaluating on a separate, unseen subset to get an honest performance estimate. You also completed a practical exercise using scikit-learn's train_test_split to split features and target, train a model, and evaluate it on the test set — a foundational skill for every data science project.

You know how to choose between a simple split and cross-validation, and you've seen how to troubleshoot common pitfalls like shape mismatches and overfitting signals.

Next up: You're ready to dive into cross-validation — a more robust way to use your data, especially when your dataset is small. Cross-validation builds on the same train/test idea but cycles through multiple splits, giving you a more stable performance estimate. That's the next lesson in this track.