Understanding Train Test Split

Train test split is a model validation procedure that reveals how your model performs on new data. Here’s how to apply it.

Written by Michael Galarnyk
Published on Jul. 28, 2022
Understanding Train Test Split
Image: Shutterstock / Built In
Brand Studio Logo

A goal of supervised learning is to build a model that performs well on new data. If you have new data, it’s a good idea to see how your model performs on it. The problem is that you may not have new data, but you can simulate this experience with a procedure like train test split. 

What Is Train Test Split?

Train test split is a model validation process that allows you to simulate how your model would perform with new data. 

This tutorial includes:

  • What is the train test split procedure?
  • How to use train test split to tune models in Python.
  • Understanding the bias-variance tradeoff.

If you would like to follow along, the code and images used in this tutorial is available on GitHub. With that, let’s get started.


What Is the Train Test Split Procedure?

Train test split is a model validation procedure that allows you to simulate how a model would perform on new/unseen data. Here is how the procedure works:

train-test-split graphic model
Train test split procedure. | Image: Michael Galarnyk


1. Arrange the Data

Make sure your data is arranged into a format acceptable for train test split. In scikit-learn, this consists of separating your full data set into “Features” and “Target.”


2. Split the Data 

Split the data set into two pieces — a training set and a testing set. This consists of random sampling without replacement about 75 percent of the rows (you can vary this) and putting them into your training set. The remaining 25 percent is put into your test set. Note that the colors in “Features” and “Target” indicate where their data will go (“X_train,” “X_test,” “y_train,” “y_test”) for a particular train test split.


3. Train the Model

Train the model on the training set. This is “X_train” and “y_train” in the image.


4. Test the Model

Test the model on the testing set (“X_test” and “y_test” in the image) and evaluate the performance.

More on Data ScienceUnderstanding Boxpolots


Consequences of Not Using Train Test Split

You could try not using train test split and instead train and test the model on the same data. However, I don’t recommend this approach as it doesn’t simulate how a model would perform on new data. It also tends to reward overly complex models that overfit on the data set.

The steps below go over how this inadvisable process works.

train test split procedure without splitting data
Train and test procedure without splitting the data. | Image: Michael Galarnyk


1. Arrange the Data

Make sure your data is arranged into a format acceptable for train test split. In scikit-learn, this consists of separating your full data set into “Features” and “Target.”


2. Train the Model

Train the model on “Features” and “Target.”


3. Test the Model

Test the model on “Features” and “Target” and evaluate the performance.

I want to emphasize again that training on an entire data set and then testing on that same data set can lead to overfitting. Overfitting is defined in the image below. The green squiggly line best follows the training data. The problem is that it is likely overfitting on the training data, meaning it is likely to perform worse on new data.

train test split overfitted data graph
Example of overfitted training data. | Image: Wikipedia

More on Data ScienceHow to Use a Z-Table and Create Your Own


Using Train Test Split In Python

train test split process for python
“Change hyperparameters” in this image is also known as hyperparameter tuning. | Image: Michael Galarnyk.

This section is about the practical application of train test split as a way to predict home prices. It spans everything from importing a data set to performing a train test split to hyperparameter tuning a decision tree regressor to predicting home prices and more.

Python has a lot of libraries that help you accomplish your data science goals including scikit-learn, pandas, and NumPy, which the code below imports.

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn import tree
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor


2. Load the Data Set

Kaggle hosts a data set which contains the price at which houses were sold for King County, which includes Seattle between May 2014 and May 2015. You can download the data set from Kaggle or load it from my GitHub. The code below loads the data set.

url = 'https://raw.githubusercontent.com/mGalarnyk/Tutorial_Data/master/King_County/kingCountyHouseData.csv'
df = pd.read_csv(url)
# Selecting columns I am interested in
columns = ['bedrooms','bathrooms','sqft_living','sqft_lot','floors','price']
df = df.loc[:, columns]
train test split home prices table
Table of home prices. | Image: Michael Galarnyk


3. Arrange Data into Features and Target

Scikit-learn’s train_test_split expects data in the form of features and target. In scikit-learn, a features matrix is a two-dimensional grid of data where rows represent samples and columns represent features. A target is what you want to predict from the data. This tutorial uses price as a target.

features = ['bedrooms','bathrooms','sqft_living','sqft_lot','floors']
X = df.loc[:, features]
y = df.loc[:, ['price']]
train test split arranged data tables
Image: Michael Galarnyk


4. Split Data Into Training and Testing Sets

train test split tables
The colors in the image indicate which variable (X_train, X_test, y_train, y_test) from the original dataframe (df) will go to for a particular train test split. If you are curious how the image was made above, I recommend you download and run the ArrangeDataKingCountySplit notebook as pandas styling functionality doesn’t always render on GitHub. | Image: Michael Galarnyk.

In the code below, train_test_split splits the data and returns a list which contains four NumPy arrays, while train_size = .75 puts 75 percent of the data into a training set and the remaining 25 percent into a testing set.

X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0, train_size = .75)

The image below shows the number of rows and columns the variables contain using the shape attribute before and after the train_test_split.

train-test-split equation
Image: Michael Galarnyk.

Shape before and after train_test_split. 75 percent of the rows went to the training set 16209/ 21613 = .75 and 25 percent went to the test set 5404 / 21613 = .25.

More on Pandas DataFrameFrom Clipboard to DataFrame With Pandas: A Quick Guide


5. What Is random_state in Train Test Split?

train-test-split random state
Image: Michael Galarnyk

The image above shows that if you select a different value for random_state, different information would go to X_train,” X_test,” y_train” and y_test”. The random_state is a pseudo-random number parameter that allows you to reproduce the same train test split each time you run the code. 

There are a number of reasons why people use random_state, including software testing, tutorials like this one and talks. However, it is recommended you remove it if you are trying to see how well a model generalizes to new data.


Train Test Split: Creating and Training a Model in Scikit-Learn

Here’s what you need to know:

4 Steps for Train Test Split Creation and Training in Scikit-Learn

  1. Import the model you want to use.
  2. Make an instance of the model.
  3. Train the model on the data.
  4. Predict labels of unseen test data. 


1. Import the Model You Want to Use

In scikit-learn, all machine learning models are implemented as Python classes.

from sklearn.tree import DecisionTreeRegressor


2. Make An Instance of the Model

In the code below, I set the hyperparameter max_depth = 2 to pre-prune my tree to make sure it doesn’t have a depth greater than two. The next section of the tutorial will go over how to choose an optimal max_depth for your tree. Also note that I made random_state = 0 so that you can get the same results as me.

reg = DecisionTreeRegressor(max_depth = 2, random_state = 0)


3. Train the Model on the Data

Train the model on the data, storing the information learned from the data.

reg.fit(X_train, y_train)


4. Predict Labels of Unseen Test Data

# Predicting multiple observations

train-test-split array

For the multiple predictions above, notice how many times some of the predictions are repeated. If you are wondering why, I encourage you to check out the code below, which will start by looking at a single observation/house and then proceed to examine how the model makes its prediction.

train test split prediction
One house’s features visualized as a Pandas DataFrame. | Image: Michael Galarnyk

The code below shows how to make a prediction for that single observation.

# predict 1 observation.

train test split array

The image below shows how the trained model makes a prediction for the one observation.

train test split price prediction chart
Image: Michael Galarnyk.

If you are curious how these sorts of diagrams are made, consider checking out my tutorial Visualizing Decision Trees using Graphviz and Matplotlib.

A train test split tutorial using Python. | Video: Michael Galarnyk


Measuring Train Test Split Model Performance

train test split formula
R² (coefficient of determination) formula. | Image: Michael Galarnyk

While there are other ways of measuring model performance such as root-mean-square error, and mean absolute error, we are going to keep this simple and use R² — known as the coefficient of determination — as our metric. 

The best possible score is 1.0. A constant model that would always predict the mean value of price would get a R² score of 0.0. However, it is possible to get a negative R² on the test set. The code below uses the trained model’s score method to return the R² of the model that we evaluated on the test set.

score = reg.score(X_test, y_test)

train test split result

You might be wondering if our R² above is good for our model. In general, the higher the R², the better the model fits the data. Determining whether a model is performing well can also depend on your field of study. Something harder to predict will generally have a lower R². My argument is that for housing data, we should have a higher R² based solely on our data.

Domain experts generally agree that one of the most important factors in housing prices is location. After all, if you are looking for a home, you’ll most likely care where it’s located. As you can see in the trained model below, the decision tree only incorporates sqft_living.

train test split flow chart
Image: Michael Galarnyk.

Even if the model was performing very well, it is unlikely that it would get buy-in from stakeholders or coworkers since there is more to a home than sqft_living.

Note that the original data set has location information like “lat” and “long.” The image below visualizes the price percentile of all the houses in the data set based on “lat” and “long,” neither were included in the data the model trained on. As you can see, there is a relationship between home price and location.

You can incorporate location information like “lat” and “long” as a way to improve the model. It’s likely places like Zillow found a way to incorporate that into their models.

train test split housing prices map
Housing price percentile for King County. | Image: Michael Galarnyk


How to Tune the max_depth of a Tree

The R² for the model trained earlier in the tutorial was about .438. However, suppose we want to improve the performance so that we can better make predictions on unseen data. While we could add more features like “lat” and “long” to the model or increase the number of rows in the data set (i.e. find more houses), we could also improve performance through hyperparameter tuning.

This involves selecting the optimal values of tuning parameters for a machine learning problem, which are often called hyperparameters. But first, we need to briefly go over the difference between parameters and hyperparameters.


Parameters vs Hyperparameters

A machine learning algorithm estimates model parameters for a given data set and updates these values as it continues to learn. You can think of a model parameter as a learned value from applying the fitting process. For example, in logistic regression you have model coefficients. In a neural network, you can think of neural network weights as a parameter. Hyperparameters or tuning parameters are metaparameters that influence the fitting process itself. 

For logistic regression, there are many hyperparameters like regularization strength C. For a neural network, there are many hyperparameters like the number of hidden layers. If all of this sounds confusing, Jason Brownlee, founder of Machine Learning Mastery, offers a good rule of thumb in his guide on parameters and hyperparameters which is: “If you have to specify a model parameter manually, then it is probably a model hyperparameter.”


Hyperparameter tuning

There are a lot of different ways to hyperparameter tune a decision tree for regression. One way is to tune the max_depth hyperparameter. The max_depth (hyperparameter) is not the same thing as depth (parameter of a decision tree), but max_depth is a way to pre-prune a decision tree. In other words, if a tree is already as pure as possible at a depth, it will not continue to split. If this isn’t clear, check out my Understanding Decision Trees for Classification (Python) tutorial to see the difference between max_depth and depth.

The code below outputs the accuracy for decision trees with different values for max_depth.

max_depth_range = list(range(1, 25))
# List to store the average RMSE for each value of max_depth:
r2_list = []
for depth in max_depth_range:
    reg = DecisionTreeRegressor(max_depth = depth,
                            random_state = 0)
    reg.fit(X_train, y_train)   
    score = reg.score(X_test, y_test)

The graph below shows that the best model R² is when the hyperparameter max_depth is equal to five. This process of selecting the best model max_depth (max_depth = 5 in this case) among many other candidate models (with different max_depth values in this case) is called model selection.

fig, ax = plt.subplots(nrows = 1, ncols = 1,
                       figsize = (10,7),
                       facecolor = 'white');
ax.set_xlim([1, max(max_depth_range)])
       axis = 'both',
       zorder = 0,
       linestyle = ':',
       color = 'k')
ax.tick_params(labelsize = 18)
ax.set_xlabel('max_depth', fontsize = 24)
ax.set_ylabel('R^2', fontsize = 24)
ax.set_title('Model Performance on Test Set', fontsize = 24)
train test split model performance graph
Image: Michael Galarnyk.

Note that the model above could have still been overfitted on the test set since the code changed max_depth repeatedly to achieve the best model. In other words, knowledge of the test set could have leaked into the model as the code iterated through 24 different values for max_depth (the length of max_depth_range is 24). This would lessen the power of our evaluation metric R², as it would no longer be as strong an indicator of generalization performance. This is why in real life, we often have training, test and validation sets when hyperparameter tuning.

More on HyperparametersRandom Forest Algorithm: A Complete Guide


Understanding the Bias-Variance Tradeoff

In order to understand why max_depth of five was the “Best Model” for our data, take a look at the graph below, which shows the model performance when tested on the training and test set.

train test split model performance graph
The code for this graph is on GitHub. | Image: Michael Galarnyk

Naturally, the training R² is always better than the test R² for every point on this graph because the models are making predictions on data they have seen before.

To the left side of the “Best Model” on the graph (anything less than max_depth = 5), we have models that underfit the data and are considered high bias because they do not not have enough complexity to learn enough about the data.

To the right side of the “Best Model” on the graph (anything more than max_depth = 5), we have models that overfit the data and are considered high variance because they are overly-complex models. These models perform well on the training data but badly on testing data.

The “Best Model” is formed by minimizing bias error — or bad assumptions in the model — and variance error — or oversensitivity to small fluctuations/noise in the training set.


Train Test Split Advantages and Disadvantages

train test split final evaluation chart
“Cross-validation: evaluating estimator performance” image from the scikit-learn documentation. | Image: scikit-learn

A goal of supervised learning is to build a model that performs well on new data, which train test split helps you simulate. With any model validation procedure it’s important to keep in mind the advantages and disadvantages. 

The advantages of train test split include: 

  • Being relatively simple and easier to understand than other methods like K-fold cross validation.
  • Avoiding overly complex models that don’t generalize well to new data.

Its disadvantages include:

  • Eliminating data that could have been used for training a machine learning model (testing data isn’t used for training).
  • Producing results that can vary for a particular train test split (random_state)
  • When hyperparameter tuning, knowledge of the test set can leak into the model. However, this can be partially solved by using a training, test, and validation set.
Hiring Now
Artificial Intelligence • Cloud • Consumer Web • Productivity • Software • App development • Data Privacy