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

Every machine-learning model needs to be tested on data it has not seen during training. Without that separation, you cannot tell whether the model has learned a real pattern or simply memorised the answers.

train_test_split is the function that creates that separation. It takes your features and targets, shuffles the rows, and returns four arrays: training features, test features, training targets, test targets. One function call replaces several lines of index slicing.

The split must happen before any preprocessing. Scaling, encoding or imputing on the full dataset and then splitting leaks information from the test rows into the training set. That leak inflates the score and hides the true error rate. The order matters more than the split ratio.

This page covers the call itself, the parameters that control the split, and the situations where a random split is not enough. Time-ordered data, grouped observations and stratified sampling each need a different approach. Knowing when the simple split applies and when it does not is half the value.

A model that has seen a row can copy its answer. That isn't a find about the model. It's a property of having been shown the answer.

So a score found on the training data measures memorisation. Some models store just right. An unconstrained choice tree will score 1.0 on its training set each time. It can keep splitting until each sample sits in its own leaf. Report that number and you're reporting the model's room to store data. This Isn't what anyone wants to know.

The question worth answering is what happens on data the model has never seen. 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 stored, 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 isn't always loud. Swapping X_test and y_train makes a shape error. This Is fine. Swapping y_train and y_test makes a program that runs and reports a meaningless number.

Any number of arrays can be passed and they're all split the same way, on the same rows. This 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 makes a different one, and so a different score.

That matters more than it sounds. If you change a setting 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're checking 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 a thing real. You can't investigate it if you can't get back to the split that failed.

Passing an int is enough. The exact value carries no meaning. random_state=0 and random_state=42 are equally any. And the only thing that matters is that it stays the same across the runs you intend to compare.

There's one honest use for leaving it out. Measuring how much the score varies between splits. This Is a real question and one that a single fixed split can't answer. The proper tool for that is cross-validation. This Repeats the exercise systematically rather than leaning on you to run it a few times.

stratify, and the split that misrepresents the data

A random split doesn't ensure 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. On a dataset where one class is rare, it can hand you a test set holding none of it at all.

stratify=y fixes this by sampling within each class. Both halves carry the same proportions as the first. The cost is nothing. For classification it should be the default rather than an option you note on hard datasets.

Where it turns key 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 found from them is noise. Stratifying ensures the proportion. This Is the min needed for the number to mean any of it.

You can stratify on a thing other than the target by passing a different array. A group label, say. Which is at times what you want when the target is smooth but some cat variable must stay weighed.

The split has to come first

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

All 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 vocab of a text vectoriser. The feature-pick choice about which columns to keep.

Scale the whole dataset and then split. The scaler has seen the test set. The test score afterwards isn't a score on unseen data. Info from those rows reached the model through the scaler's params. The effect is often small and at times huge, and it's always in the direction of making the model look better than it's.

The rule that follows is short: split first. Fit each transformer on the training half only. Pipelines exist largely to make that structurally hopeless to get wrong. This Is why they arrive later in this track and why they aren't optional in real work.

How big should the test set be

The rule is 20% or 25%. The rule is a compromise between two things pulling in opposite directions.

A larger test set gives a more solid guess. The score is found from more samples and is less at the mercy of which given rows landed there. A smaller test set leaves more data for training. This Often makes a better model.

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

The other consideration is the rarest class. A test set that holds four examples of a thing can only report fit on it in steps of 25%. No amount of careful metric choice recovers from that.

Time changes the rules

When rows are sorted in time, shuffling is wrong.

Shuffling puts rows from after the test period into the training set. The model learns from the future and is then asked to predict the past. It will do well. The score will be worthless. The situation it was scored in can't occur in use. In production, the future is just what you don't have.

shuffle=False keeps the order and takes the last portion as the test set. This Is the right shape. Train on the past, test on the more recent. Note that stratify can't be used with shuffle=False. The two are conceptually incompatible anyway.

For any of it more careful, TimeSeriesSplit gives the cross-validation equal. A series of splits, each training on all up to a point and testing on what comes next. The same thinking applies to any structure with groups that must not be broken across the split. A few rows per patient, per user, per doc. Where GroupShuffleSplit keeps a group fully 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 hasn't seen. It stops answering that question the moment you use it to make a choice.

This is the part that gets lost. If you fit five models, look at the test score for each. Pick the best, you've used the test set to choose a model. And the winner's score is now optimistic. It was selected for doing well on those given rows. Do it a dozen times, tweaking as you go. 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 settings. test touched just once at the very end to report a number. the validation half is often replaced by cross-validation on the training data. This Uses the data better. The test set is still set aside and left alone.

The discipline is easier to state than to keep. Each look at the test set costs a little of its honesty. Nested cross-validation exists for the situation where you can't afford even that. 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 free, and often they aren't.

a few rows per entity. Ten readings from the same patient, a few photos 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. The model can recognise the entity rather than learn the pattern. 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's so easy to get wrong: shuffle=False, or TimeSeriesSplit.

Nested or hierarchical data. Pupils within schools, readings within sites. The same thinking as groups. If the model can spot the container, it will.

The question to ask before splitting is simply. Could two rows on opposite sides of the split share a thing 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's a rule 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 any of it that learns from the data, yes. Dropping obviously corrupt rows is fine either way. Imputing a median isn't.

Things to try with Splitting Train and Test

  1. Run the fourth editor. The training score is 1.0 and the test score isn't. That gap is the entire reason this page exists.
  2. Drop 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. 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 holds rows from the beginning of the series.

Questions people ask about Splitting Train and Test

What test size should I use? A common choice is 20 to 30 percent. Smaller datasets need a larger share held out to get a stable estimate, while huge ones can spare far less.

What does random_state do in train_test_split? It fixes the shuffle so the same rows land in the same halves on every run. Without it, scores change between runs.

When should I use stratify? When classes are imbalanced or the dataset is small. stratify=y keeps the class proportions the same in both halves.

Can I use train_test_split for time series? Not with shuffling. Time-ordered data needs the past to train and the future to test, for example with shuffle=False or TimeSeriesSplit.

Should I scale before or after splitting? After. Split first, fit the scaler on the training part only, then transform both parts.

Is one split enough to judge a model? It gives one estimate that depends on which rows landed where. Cross-validation averages over several splits and is steadier.

Where this leaves you

One line, four objects, and three args worth setting on purpose each time: random_state so the test is repeatable, stratify so the halves resemble each other, and shuffle=False when the rows are sorted in time. What it gives you is a number you're 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

Split data into training and test sets with train_test_split, and see why scoring a model on the data it memorised tells you nothing.

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.