Lesson menuPython TutorialHideShow

Python Tutorial

Train / Test Split

Measure how well a model generalizes by evaluating it on data it never saw during training.

← Feature ScalingDecision Tree →

Why Split the Data?

A model can memorize its training data and still fail on new inputs — a problem called overfitting. To estimate real-world performance you hold back a portion of the data as a test set, train on the rest, and score the model on the untouched test set.

A common split is 80% training and 20% testing. With very large datasets a smaller test fraction (e.g. 1–5%) still gives a reliable estimate.

Splitting with scikit-learn

import numpy as np
from sklearn.model_selection import train_test_split

rng = np.random.default_rng(2)
x = rng.normal(3, 1, 100)
y = rng.normal(150, 40, 100) / x

X = x.reshape(-1, 1)   # sklearn expects 2-D features

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42)

print(X_train.shape, X_test.shape)   # (80, 1) (20, 1)

random_state fixes the shuffling so your split is reproducible. Use the same seed to compare experiments fairly.

Fit on Train, Score on Test

from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import make_pipeline
from sklearn.metrics import r2_score

model = make_pipeline(PolynomialFeatures(4), LinearRegression())
model.fit(X_train, y_train)

print("Train R2:", r2_score(y_train, model.predict(X_train)).round(3))
print("Test  R2:", r2_score(y_test,  model.predict(X_test)).round(3))

If the training score is high but the test score is much lower, the model is overfitting. If both are low, it is underfitting.

Stratified Splits for Classification

When classes are imbalanced, a random split can leave a class under-represented in the test set. stratify keeps the class proportions identical in both parts.

X_train, X_test, y_train, y_test = train_test_split(
    X, labels, test_size=0.2, random_state=42, stratify=labels)

Train / Validation / Test

For serious model tuning use three sets: train to fit, validation to tune hyperparameters, and test for a final, one-time estimate. Split twice to build them.

# First carve off the test set, then split the rest into train/val
X_temp, X_test, y_temp, y_test = train_test_split(
    X, y, test_size=0.15, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(
    X_temp, y_temp, test_size=0.1765, random_state=42)  # ~0.15 of total

Touch the test set only once, at the very end. Repeatedly tuning against it leaks information and produces over-optimistic scores.

Best Practices

  • Split before any preprocessing, or wrap preprocessing in a Pipeline to avoid leakage.
  • Stratify classification splits so rare classes appear in every set.
  • Compare train vs test scores to diagnose over/underfitting.
  • For small datasets, prefer cross-validation over a single split for a more stable estimate.

Try It Yourself

Exercise 1: Split 100 samples into 70% train / 30% test and print both sizes.

Show solution
import numpy as np
from sklearn.model_selection import train_test_split
X = np.arange(100).reshape(-1, 1)
tr, te = train_test_split(X, test_size=0.3, random_state=0)
print(len(tr), len(te))   # 70 30

Exercise 2: A model scores 0.98 on train but 0.62 on test. What is happening?

Show solution

Overfitting — the model memorized the training data and fails to generalize. Simplify it, add data, or regularize.

📘 Real-World Deep Dive

Knowing <strong>ML Train Test (scikit-learn)</strong> well is what turns scikit-learn from a curiosity into a daily tool — you'll reach for it in nearly every real project.

Real-Life Scenario

An end-to-end usage of ML Train Test that you'd actually see in a data pipeline or analytics notebook.

Real-Life Example

import numpy as np
from sklearn.model_selection import train_test_split
X = np.arange(20).reshape(10, 2)
y = np.arange(10)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.25, random_state=0)
print("train:", len(Xtr), "test:", len(Xte))

Expected Output

(see source)

Common mistakes

  • fit expects numeric arrays; OneHotEncoder / LabelEncoder are easy to forget for categorical features.
  • Calling predict on a model trained on unscaled data and then scaling inputs at inference time silently degrades accuracy.
  • train_test_split(X, y) requires both arrays; train_test_split(X) for unsupervised learning slips past static checkers.
  • Treating ML Train Test as a black box without reading the docs — the API has subtle defaults that bite when you scale.

🚀 Performance & Best Practices

  • Wrap preprocessing + estimator in a Pipeline so fit / predict stay reproducible.
  • Use joblib / pickle for serialising models, not the entire Python state.
  • For > 100 k rows, switch to HistGradientBoostingClassifier or cuML.
  • When working with scikit-learn, prefer vectorised / batched operations over Python loops.

🧪 Try It Yourself

  1. Reproduce the snippet on a representative slice of your own data.
  2. Profile the snippet with cProfile or timeit and find the single biggest improvement.
  3. Generalise the snippet into a small, reusable function you can drop into future projects.

FAQ: Train / Test Split

Common questions about this page.

What is Train / Test Split?

Train / Test Split is a Machine Learning lesson that explains train / test split in Python. Measure how well a model generalizes by evaluating it on data it never saw during training. Copy the samples and run them in the Python editor. It is written for beginners who want a clear definition and working examples.

Should I run train / test split examples locally for better learning?

Yes. Use the browser editor on StudyGrid for a quick check, then Download the example and run it on your computer. Local runs show real errors and the real toolchain, which is one of the fastest ways to learn train / test split in this Machine Learning Python lesson (Train / Test Split).

How do I use train / test split in Python?

To use train / test split in Python, follow the examples on this StudyGrid page. Copy a snippet, run it in the browser, then Download and run it locally for better learning. Change the values and compare the output.

What is the syntax of train / test split?

This Train / Test Split tutorial shows train / test split syntax with short Python examples. Use the code blocks in this lesson for the exact statements, then try them in your editor.

Train / Test Split example for beginners

Yes. This page includes a beginner train / test split example you can copy and run. It is designed for searches such as "train / test split for beginners", "train / test split example", and "how to use train / test split".

What are common mistakes with train / test split?

Common train / test split mistakes include wrong syntax, mixing types, and skipping practice. Work through this Machine Learning chapter in order, run every example, and check the output before moving on.

Why should I learn train / test split?

Train / Test Split is used in real Python work. Learning train / test split helps you write clearer programs and continue the Machine Learning tutorial on StudyGrid.

Is Train / Test Split free to learn online?

Yes. You can learn train / test split free on StudyGrid (studygrid.in). This chapter is part of the Machine Learning path and includes examples, syntax, and next-step links.