Skip to content

Split Data into Train and Test Data

After training a model, you should evaluate its performance and get an estimate of its accuracy when applied to new data. You cannot use the train data again for this, since this will give you an overestimate of its performance on new data. However, you can solve this by training the model on a portion of the data (~70%) and evaluate its performance on the rest. These are called the training set and test set respectively. This recipe tells you how you can do this using the scikit-learn's train_test_split() function.

# Load packages
import numpy as np 
import pandas as pd 
from sklearn.model_selection import train_test_split
%config InlineBackend.figure_format = 'retina'
# Upload your data as CSV and load as data frame
df = pd.read_csv('housing.csv')
df.head()

The dataset (source: Kaggle) consists of information about houses in California. We will now split this dataset into four parts using the library. These four parts are: X_train, X_test, y_train and y_test.

TO_PREDICT = 'median_house_value'                                                  # The column that we would like to predict; the output
TEST_SIZE = .30                                                                    # The proportion of data that will be used in the tests

def split_in_X_y(df):                                                              # This function first splits the dataset into the output (y) and the input (X) of the model
    X = df.loc[:, df.columns != TO_PREDICT]
    y = df[TO_PREDICT]
    
    return X,y

X,y = split_in_X_y(df)

X_train, X_test, y_train, y_test = train_test_split(X,y,test_size = TEST_SIZE)     # Finally, the data is split into train and test data using the scikit-learn package

And that's it. The data is now split into the four parts. Let's save these dataframes into four seperate csv-files.

X_train.to_csv('X_train.csv')
y_train.to_csv('y_train.csv')
X_test.to_csv('X_test.csv')
y_test.to_csv('y_test.csv')