Skip to main content
MahiDev — Learn. Build. Become AI-ready.

Lesson 03

Training and Test Sets

Hold back rows the model never sees with train_test_split, repeat the split with random_state, keep class balance with stratify, and score honestly.

Intermediate25 min
Training and Test Sets hero Module 27 Lesson 3Training and Test Sets hero Module 27 Lesson 3

A model that is scored on the same rows it learned from can look perfect while being useless on new data, like a learner who memorised last year's answers. The fix is to hide some rows before training. The model learns from the training set and is scored on the test set, which stands in for the future.

Rows shuffled and split into 160 training rows and 40 test rowsRows shuffled and split into 160 training rows and 40 test rows

Why hold rows back

A deep decision tree can memorise every training row. Scored on those same rows it looks excellent, which tells you nothing about a learner it has not met.

python
import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

from sklearn.tree import DecisionTreeRegressor

X = learners[["hours", "practice"]]
y = learners["score"]
tree = DecisionTreeRegressor(random_state=0).fit(X, y)
print("score on the rows it learned:", round(tree.score(X, y), 3))

Output:

text
score on the rows it learned: 0.999

Split with train_test_split

train_test_split shuffles the rows and returns four pieces in a fixed order: training features, test features, training labels, test labels. test_size=0.2 keeps a fifth of the rows for testing.

python
import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

from sklearn.model_selection import train_test_split

X = learners[["hours", "practice"]]
y = learners["passed"]
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)
print(y_train.shape, y_test.shape)

Output:

text
(160, 2) (40, 2)
(160,) (40,)

random_state makes the split repeatable

The shuffle is random, so each run could pick different test rows. A fixed random_state makes it the same every time, which keeps results comparable while you experiment.

python
import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

from sklearn.model_selection import train_test_split

X = learners[["hours", "practice"]]
first = train_test_split(X, test_size=0.2, random_state=42)[1]
again = train_test_split(X, test_size=0.2, random_state=42)[1]
other = train_test_split(X, test_size=0.2, random_state=7)[1]

print(first.index[:5].tolist())
print(first.index.equals(again.index), first.index.equals(other.index))

Output:

text
[95, 15, 30, 158, 128]
True False

stratify keeps the class balance

48 percent of learners passed. A random split can land a test set with noticeably more or fewer passes. stratify=y keeps the same share in both parts, which matters most for small or unbalanced data.

python
import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

from sklearn.model_selection import train_test_split

X = learners[["hours", "practice"]]
y = learners["passed"]
_, _, _, plain = train_test_split(X, y, test_size=0.2, random_state=3)
_, _, _, even = train_test_split(X, y, test_size=0.2, random_state=3, stratify=y)

print("whole table:", y.mean())
print("plain split test:", plain.mean())
print("stratified test:", even.mean())

Output:

text
whole table: 0.48
plain split test: 0.625
stratified test: 0.475

An honest score

Fit on the training rows only, then score both parts. The gap between the two numbers is the part of the training score that was memorisation.

python
import numpy as np
import pandas as pd

rng = np.random.default_rng(27)
learners = pd.DataFrame({
    "hours": rng.uniform(1, 12, 200).round(1),
    "practice": rng.integers(0, 41, 200),
    "video_minutes": rng.integers(0, 3001, 200),
    "track": rng.choice(["Python", "Data", "Web"], 200),
})
bonus = learners["track"].map({"Python": 0, "Data": 6, "Web": -5})
noise = rng.normal(0, 6, 200)
learners["score"] = (
    22 + learners["hours"] * 3.4 + learners["practice"] * 0.55 + bonus + noise
).clip(0, 100).round()
learners["passed"] = (learners["score"] >= 60).astype(int)

from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor

X = learners[["hours", "practice"]]
y = learners["score"]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
tree = DecisionTreeRegressor(random_state=0).fit(X_train, y_train)

print("train:", round(tree.score(X_train, y_train), 3))
print("test: ", round(tree.score(X_test, y_test), 3))

Output:

text
train: 0.999
test:  0.442

Key takeaways

  • Never judge a model on the rows it was trained on

  • train_test_split returns X_train, X_test, y_train, y_test in that order

  • random_state fixes the shuffle so results are repeatable

  • stratify=y keeps the class balance the same in both parts

  • A big train and test gap means the model memorised instead of learning

  • Next: Linear Regression

  • Previous: Features and Labels

Loading playground…

Check your understanding

Training and Test Sets Quiz

Why keep a test set?

1 / 10

Previous

← Features and Labels

Next

Linear Regression →