Machine Learning: Training and Testing in Python

The train/test split is the most fundamental step in building a machine learning model. The idea is simple: keep some of your data hidden from the model during training, then measure how well the model performs on that hidden portion. Without this separation, you have no honest way to know whether your model has genuinely learned a pattern or merely memorised the training examples.

This chapter covers how train_test_split from scikit-learn works, what each parameter does, and how to evaluate the resulting model — for both regression and classification problems.

Why Separate Training from Testing Data?

A model that trains and evaluates on the same data will look far more accurate than it really is. This is called data leakage or overfitting: the model has memorised the training examples rather than learning a generalizable pattern.

Consider a student who studies 100 practice questions and then takes a test using those same 100 questions. They may score 100% — but that score tells you nothing about whether they understand the subject.

The train/test split prevents this by:

• Training the model on one portion of the data so it can learn patterns.
• Testing the model on a separate, unseen portion to measure real-world performance.

A typical split is 80% training / 20% testing, though the right ratio depends on how much data you have.

The train_test_split Function

scikit-learn provides train_test_split in its model_selection module. It randomly shuffles your dataset and splits it into two parts:

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
)

The four return values are always in this order: training features, testing features, training labels, testing labels.

Key Parameters

Effect of 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)}")

Output:

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 and Reproducibility

Without a random_state, the split is different every time you run the script. Set it to any integer to get reproducible results:

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

The choice of integer (42, 0, 1, etc.) does not matter — as long as you use the same value consistently.

Regression Example: Predicting House Prices

The following example generates a synthetic dataset of house sizes and prices, trains a linear regression model, and evaluates it on the held-out test set.

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}")

Output:

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

Interpreting the metrics:

• MSE (Mean Squared Error) is the average squared difference between predictions and actual values. Lower is better, but the scale depends on your target variable (here, dollars).
• R² ranges from 0 to 1. A value of 0.965 means the model explains about 96.5% of the variance in house prices — a strong fit on this simple dataset.

For more on linear regression, see Linear Regression in Python.

Classification Example: Iris Flower Species

For a classification task, accuracy and a classification report are more informative than 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"]))

Output:

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

Why stratify=y matters: Without it, a random split of an imbalanced dataset could put most samples of a rare class into training, leaving none in the test set. stratify=y guarantees each class appears in the same proportions in both splits.

For more on classification, see Logistic Regression in Python and Confusion Matrix in Python.

Common Gotchas

Preprocessing Before or After the Split?

Always split before fitting any preprocessing steps (scaling, encoding, imputation). If you scale all your data and then split, the test set has influenced the scaler — a form of data leakage.

The correct order:

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

See Feature Scaling in Python for a full walkthrough.

Shuffling and Time-Series Data

train_test_split shuffles data by default. For time-series data this is wrong — you would be training on future data to predict the past. Set shuffle=False and make sure your data is sorted chronologically before splitting:

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

When a Single Split Is Not Enough

A single 80/20 split gives you one estimate of model performance that depends on which samples happened to land in the test set. Cross-validation repeats the split multiple times and averages the scores, giving a much more stable estimate — especially on small datasets.

Choosing an Evaluation Metric

The right metric depends on your problem type:

For binary classification, see AUC-ROC Curve in Python and Confusion Matrix in Python. For hyperparameter tuning after you have a working split, see Grid Search in Python.

Summary

• Split your data before any preprocessing to avoid data leakage.
• Use test_size=0.2 as a reasonable default; adjust based on dataset size.
• Set random_state to any integer to get reproducible splits.
• Use stratify=y for classification tasks, especially on imbalanced data.
• Set shuffle=False for time-series data.
• A single train/test split is a quick baseline; use cross-validation for more reliable performance estimates.