Splitting Train and Test

The one line that separates a score you can report from a number that means nothing at all.

Overview

The problem it solves

A model that has seen a row can reproduce its answer. That is not a discovery about the model; it is a property of having been shown the answer.

So a score computed on the training data measures memorisation, and some models memorise perfectly. An unconstrained decision tree will score 1.0 on its training set essentially every time, because it can keep splitting until every sample sits in its own leaf. Report that number and you are reporting the model's capacity to store data, which is not what anyone wants to know.

The question worth answering is what happens on data the model has never seen, because that is the only situation it will ever face in use. Holding some rows back and never letting the model touch them is the simplest honest way to find out.

Worth knowing

The four returns are X_train, X_test, y_train, y_test - both X pieces before either y piece.
random_state makes the split reproducible; without it two runs are two different experiments.
stratify=y keeps each class in the same proportion in both halves, and matters most when a class is rare.
A score measured on the training data is not a score - the model has already seen every one of those rows.
test_size takes a fraction as a float or a row count as an integer.
shuffle=False for anything ordered in time, or the model is trained on the future.

Splitting Train and Test: A Practical Guide

A model's score on the data it learned from tells you how well it memorised, not how well it works. The split is what turns one into the other.

The split, and the order of the four returns

One call, four objects, and an order that catches everyone the first time.

example_01.pyscikit-learn
Output

random_state makes the split reproducible

Without it, every run is a different experiment and two scores cannot be compared.

example_02.pyscikit-learn
Output

stratify keeps the class balance

A random split can hand you a test set whose class proportions differ from the data it came from.

example_03.pyscikit-learn
Output

Why the training score is not a score

The same model, measured twice. Only one of the two numbers is evidence of anything.

example_04.pyscikit-learn
Output

test_size takes a fraction or a count

A float is a proportion, an integer is a number of rows.

example_05.pyscikit-learn
Output

Time series must not be shuffled

Shuffling puts future rows in the training set, which is a way of scoring well and learning nothing.

example_06.pyscikit-learn
Output

What the call does

train_test_split shuffles the rows and cuts them into two groups, returning four objects: the features and target for training, then the features and target for testing.

The return order is X_train, X_test, y_train, y_test — both X pieces before either y piece — and getting it wrong is the single most common early mistake with this function. The failure is not always loud: swapping X_test and y_train produces a shape error, which is fine, but swapping y_train and y_test produces a program that runs and reports a meaningless number.

Any number of arrays can be passed and they are all split the same way, on the same rows, which is what keeps X and y aligned. That property is the whole reason to use the function rather than slicing by hand.

random_state, and what reproducibility buys

The split is random, so without a seed each run produces a different one, and therefore a different score.

That matters more than it sounds. If you change a hyperparameter and the score moves from 0.94 to 0.96, you need to know whether the change did that or whether the split did. With a fixed seed, the split is held constant and the difference is attributable. Without one, you are comparing two things that differ in two ways.

It also makes a bug reproducible. A model that fails on one split and not another is telling you something real, and you cannot investigate it if you cannot get back to the split that failed.

Passing an integer is enough. The specific value carries no meaning — random_state=0 and random_state=42 are equally arbitrary — and the only thing that matters is that it stays the same across the runs you intend to compare.

There is one honest use for leaving it out: measuring how much the score varies between splits, which is a real question and one that a single fixed split cannot answer. The proper tool for that is cross-validation, which repeats the exercise systematically rather than relying on you to run it a few times.

stratify, and the split that misrepresents the data

A random split does not guarantee the two halves look alike. On a dataset with three equal classes, a plain split can easily hand you a test set with 17 of one and 14 of another, and on a dataset where one class is rare, it can hand you a test set containing none of it at all.

stratify=y fixes this by sampling within each class, so both halves carry the same proportions as the original. The cost is nothing, and for classification it should be the default rather than an option you remember on difficult datasets.

Where it becomes essential is imbalance. With 1% positives and a 20% test set, an unstratified split has a real chance of putting so few positives in the test set that the metric computed from them is noise. Stratifying guarantees the proportion, which is the minimum needed for the number to mean anything.

You can stratify on something other than the target by passing a different array — a group label, say — which is occasionally what you want when the target is continuous but some categorical variable must stay balanced.

The split has to come first

The order of operations matters, and getting it wrong is the most common way to produce an inflated score.

Everything learned from data must be learned from the training set alone. That includes the obvious — the model — and the less obvious: the mean and standard deviation used for scaling, the categories known to an encoder, the median used to fill missing values, the vocabulary of a text vectoriser, and the feature-selection decision about which columns to keep.

Scale the whole dataset and then split, and the scaler has seen the test set. The test score afterwards is not a score on unseen data, because information from those rows reached the model through the scaler's parameters. The effect is usually small and occasionally enormous, and it is always in the direction of making the model look better than it is.

The rule that follows is short: split first, and fit every transformer on the training half only. Pipelines exist largely to make that structurally impossible to get wrong, which is why they arrive later in this track and why they are not optional in real work.

How big should the test set be

The convention is 20% or 25%, and the convention is a compromise between two things pulling in opposite directions.

A larger test set gives a more reliable estimate, because the score is computed from more samples and is less at the mercy of which particular rows landed there. A smaller test set leaves more data for training, which usually produces a better model.

With a lot of data, the tension disappears: 1% of a million rows is ten thousand test samples, which is plenty. With a few hundred rows, both halves are uncomfortable — the estimate is noisy and the model is starved — and that is precisely the situation where a single split should be replaced by cross-validation, which uses every row for both purposes without ever training and testing on the same one.

The other consideration is the rarest class. A test set that contains four examples of something can only report accuracy on it in steps of 25%, and no amount of careful metric choice recovers from that.

Time changes the rules

When rows are ordered in time, shuffling is wrong.

Shuffling puts rows from after the test period into the training set, so the model learns from the future and is then asked to predict the past. It will do well, and the score will be worthless, because the situation it was scored in cannot occur in use — in production, the future is exactly what you do not have.

shuffle=False keeps the order and takes the last portion as the test set, which is the right shape: train on the past, test on the more recent. Note that stratify cannot be used with shuffle=False, and the two are conceptually incompatible anyway.

For anything more careful, TimeSeriesSplit provides the cross-validation equivalent: a series of splits, each training on everything up to a point and testing on what comes next. The same reasoning applies to any structure with groups that must not be broken across the split — several rows per patient, per user, per document — where GroupShuffleSplit keeps a group entirely on one side.

What the test set is for, and what it is not

A held-out set answers one question: how does this model behave on data it has not seen. It stops answering that question the moment you use it to make a decision.

This is the part that gets lost. If you fit five models, look at the test score for each, and pick the best, you have used the test set to choose a model — and the winner's score is now optimistic, because it was selected for doing well on those particular rows. Do it a dozen times, tweaking as you go, and the test set has been fitted to as surely as if you had trained on it, just more slowly and by hand.

The standard remedy is three sets rather than two. Train to fit the model, validation to compare models and tune hyperparameters, and test touched exactly once at the very end to report a number. In practice the validation half is usually replaced by cross-validation on the training data, which uses the data better, and the test set is still set aside and left alone.

The discipline is easier to state than to keep: every look at the test set costs a little of its honesty. Nested cross-validation exists for the situation where you cannot afford even that, and is the correct answer when the difference between two models is small enough to matter.

Splits that respect structure

A plain random split assumes rows are independent, and often they are not.

Several rows per entity. Ten readings from the same patient, several photographs of the same object, multiple purchases by one customer. A random split puts some of an entity's rows in training and others in test, so the model can recognise the entity rather than learn the pattern, and the score is inflated by an amount nothing in the output reveals. GroupShuffleSplit and GroupKFold take a groups array and keep each group whole.

Time. Covered above, and worth repeating because it is so easy to get wrong: shuffle=False, or TimeSeriesSplit.

Nested or hierarchical data. Pupils within schools, measurements within sites. The same reasoning as groups — if the model can identify the container, it will.

The question to ask before splitting is simply: could two rows on opposite sides of the split share something that would let the model cheat? If the answer is yes, a plain random split will overstate the score, and the amount is unpredictable.

Is 80/20 a rule? It is a convention that suits a few thousand rows. With a million, 1% is a fine test set; with two hundred, a single split is the wrong tool and cross-validation is the right one.

Should the test set be split off before cleaning? Before anything that learns from the data, yes. Dropping obviously corrupt rows is fine either way; imputing a median is not.

Things to try

  1. Run the fourth editor. The training score is 1.0 and the test score is not. That gap is the entire reason this page exists.
  2. Remove random_state. Run the same editor three times and watch the test score move. Then put it back.
  3. Make a class rare. Keep only ten samples of class 2, split without stratify, and count the classes in the test set a few times with different seeds.
  4. Break the order. In the last editor, compare the shuffled and unshuffled test rows. The shuffled set contains rows from the beginning of the series.

Where this leaves you

One line, four objects, and three arguments worth setting deliberately every time: random_state so the experiment is repeatable, stratify so the halves resemble each other, and shuffle=False when the rows are ordered in time. What it gives you is a number you are entitled to report.

Check yourself

0 of 4

Answer without scrolling back up.

  1. What order does train_test_split return its four objects in?

  2. Why pass random_state?

  3. What does stratify=y do?

  4. Why must scaling happen after the split?

Cheat sheet

Splitting Train and Test

A model's score on the data it learned from tells you how well it memorised, not how well it works. The split is what turns one into the other.

SCIKIT-LEARN · vizlearn.in/sklearn/train_test_split.html

About the author

Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.