Skip to content

Creating a Test Set (Avoiding Data Snooping)

What you’ll learn

  • what data snooping bias is and why your brain is the problem
  • purely random sampling with train_test_splittrain_test_split
  • why random sampling can introduce sampling bias on small datasets
  • stratified sampling with StratifiedShuffleSplitStratifiedShuffleSplit
  • how to compare stratified vs. random test sets

Why set aside data you haven’t even explored yet?

It sounds backwards: you’ve barely glanced at the data, and already you’re told to lock a chunk of it away. But your brain is an excellent pattern-detection machine — which means it’s also excellent at overfitting. If you explore the full dataset first, you’re likely to spot a pattern that only exists by chance, pick a model or feature because of it, and then get an overly optimistic evaluation when you finally test on that same data. This is called data snooping bias, and it’s one of the easiest ways to accidentally fool yourself into shipping a model that performs worse in the real world than your numbers suggest.

The fix is simple in principle: split off a test set immediately, put it aside, and don’t look at it again until you’re evaluating your final model.

diagram Train/test split, once and for all mermaid
Split before exploring, and never touch the test set again until the final evaluation.

Purely random sampling

The simplest approach: shuffle the rows and take a random ~20% as your test set. Scikit-learn’s train_test_split()train_test_split() does this for you, with a random_staterandom_state so the split is reproducible:

Random train/test split
from sklearn.model_selection import train_test_split
 
train_set, test_set = train_test_split(housing, test_size=0.2, random_state=42)
 
print(len(train_set), len(test_set))
Random train/test split
from sklearn.model_selection import train_test_split
 
train_set, test_set = train_test_split(housing, test_size=0.2, random_state=42)
 
print(len(train_set), len(test_set))

Setting random_state=42random_state=42 guarantees you get the same split every time you rerun the notebook — without it, a fresh random split each run means your model eventually sees every row, which defeats the purpose of a test set.

Keeping the split stable as data changes

A fixed random_staterandom_state is reproducible today, but it isn’t stable if the dataset ever grows. Fetch an updated CSV tomorrow with 500 new rows appended, and the row positions train_test_splittrain_test_split uses internally no longer line up the same way — an instance that used to sit in the training set can jump into the test set, and vice versa. Your model would quietly start training on rows it was supposed to be evaluated against later.

The fix is to base the split on each instance’s own identifier, not its row number, so the same row always lands in the same set no matter how many new rows show up afterward. Hashing the identifier gives you a threshold rule that stays consistent as long as new rows are only ever appended, never re-ordered or deleted:

Identifier-based (hash) split
from zlib import crc32
import numpy as np
 
def is_id_in_test_set(identifier, test_ratio):
    return crc32(np.int64(identifier)) & 0xffffffff < test_ratio * 2**32
 
def split_train_test_by_id(data, test_ratio, id_column):
    ids = data[id_column]
    in_test_set = ids.apply(lambda id_: is_id_in_test_set(id_, test_ratio))
    return data.loc[~in_test_set], data.loc[in_test_set]
Identifier-based (hash) split
from zlib import crc32
import numpy as np
 
def is_id_in_test_set(identifier, test_ratio):
    return crc32(np.int64(identifier)) & 0xffffffff < test_ratio * 2**32
 
def split_train_test_by_id(data, test_ratio, id_column):
    ids = data[id_column]
    in_test_set = ids.apply(lambda id_: is_id_in_test_set(id_, test_ratio))
    return data.loc[~in_test_set], data.loc[in_test_set]

The housing dataset doesn’t ship with a ready-made identifier column, so you build one — the row index works if you only ever append new rows and never delete old ones, or you can combine stable columns (ones guaranteed not to change) into a synthetic ID:

Build a stable identifier and split
housing_with_id = housing.reset_index()  # adds an "index" column
train_set, test_set = split_train_test_by_id(housing_with_id, 0.2, "index")
 
# more robust: a synthetic ID from stable, never-changing columns
housing_with_id["id"] = housing["longitude"] * 1000 + housing["latitude"]
train_set, test_set = split_train_test_by_id(housing_with_id, 0.2, "id")
Build a stable identifier and split
housing_with_id = housing.reset_index()  # adds an "index" column
train_set, test_set = split_train_test_by_id(housing_with_id, 0.2, "index")
 
# more robust: a synthetic ID from stable, never-changing columns
housing_with_id["id"] = housing["longitude"] * 1000 + housing["latitude"]
train_set, test_set = split_train_test_by_id(housing_with_id, 0.2, "id")

For a static, one-off dataset like the one used in this phase, plain train_test_split(..., random_state=42)train_test_split(..., random_state=42) is perfectly fine — it’s worth knowing why a production system that keeps retraining on a growing dataset needs the identifier-based version instead.

Why purely random sampling can go wrong

Purely random sampling works well when your dataset is large relative to the number of attributes. But on a smaller dataset — or when one feature is especially important — a random split can introduce sampling bias. It’s the same problem a survey company faces: if the US population is 51.3% female and you call 1,000 random people, there’s roughly a 12% chance your sample skews below 49% or above 54% female, even though you sampled “randomly.”

Suppose domain experts tell you median_incomemedian_income is the strongest predictor of house prices. You’d want your test set to have the same proportions of income brackets as the full dataset — that’s stratified sampling: divide the population into homogeneous groups (“strata”) and sample the correct number from each stratum.

First, turn the continuous median_incomemedian_income into discrete income categories:

Create income categories
import numpy as np
import pandas as pd
 
housing["income_cat"] = pd.cut(
    housing["median_income"],
    bins=[0.0, 1.5, 3.0, 4.5, 6.0, np.inf],
    labels=[1, 2, 3, 4, 5],
)
Create income categories
import numpy as np
import pandas as pd
 
housing["income_cat"] = pd.cut(
    housing["median_income"],
    bins=[0.0, 1.5, 3.0, 4.5, 6.0, np.inf],
    labels=[1, 2, 3, 4, 5],
)

Then split with StratifiedShuffleSplitStratifiedShuffleSplit, using income_catincome_cat as the stratum:

Stratified train/test split
from sklearn.model_selection import StratifiedShuffleSplit
 
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
 
for train_index, test_index in split.split(housing, housing["income_cat"]):
    strat_train_set = housing.loc[train_index]
    strat_test_set = housing.loc[test_index]
 
# income_cat was only needed for stratifying — drop it afterwards
for set_ in (strat_train_set, strat_test_set):
    set_.drop("income_cat", axis=1, inplace=True)
Stratified train/test split
from sklearn.model_selection import StratifiedShuffleSplit
 
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
 
for train_index, test_index in split.split(housing, housing["income_cat"]):
    strat_train_set = housing.loc[train_index]
    strat_test_set = housing.loc[test_index]
 
# income_cat was only needed for stratifying — drop it afterwards
for set_ in (strat_train_set, strat_test_set):
    set_.drop("income_cat", axis=1, inplace=True)

Compare the income category proportions across the full dataset, the stratified test set, and a purely random test set, and you’ll see the stratified split lines up almost exactly with the full dataset, while the random split can be noticeably skewed.

Visualize it

sketch Stratified vs. random test split p5.js
Same 5 income strata, sampled two ways - stratified tracks the full dataset as the bars grow in; a purely random split keeps drifting to a different (worse) mix each time.

🧪 Try It Yourself

Exercise 1 – Split with train_test_split

Exercise 2 – Bucket a continuous column with pd.cut

Exercise 3 – Stratified sampling proportions

Next

Continue to Exploratory Data Analysis & Correlations — now that the test set is safely locked away, it’s time to dig into the training data.

If this helped you, consider buying me a coffee ☕

Buy me a coffee

Was this page helpful?

Let us know how we did