A Beginner’s Guide to Splitting Machine Learning Datasets in Python

Splitting a machine learning dataset into training and test sets is an important step in the model-building process, as it allows us to evaluate the model’s performance on unseen data. In this tutorial, we will learn how to split a dataset into training and test sets in Python using the scikit-learn library.

1. Importing libraries and dataset

First, we need to import the necessary libraries and the dataset that we will be using for this tutorial. We will be using the scikit-learn library for splitting the dataset, and the pandas library for reading and manipulating the dataset.

1
import pandas as pd
2
from sklearn.model_selection import train_test_split

Next, we will read the dataset into a Pandas DataFrame using the read_csv function.

1
df = pd.read_csv('my_dataset.csv')

2. Splitting the dataset

To split the dataset into training and test sets, we can use the train_test_split function from the scikit-learn library. This function shuffles the rows of the dataset and splits it into two datasets: a training set, which is used to train the model, and a test set, which is used to evaluate the model’s performance.

1
X = df.drop('target', axis=1)
2
y = df['target']
3

4
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

The test_size parameter specifies the proportion of the dataset that should be included in the test set. For example, a test_size of 0.2 means that 20% of the data will be included in the test set and 80% will be included in the training set.

3. Stratified sampling

In some cases, it may be important to ensure that the training and test sets have a similar class distribution to the original dataset. This is particularly important if the dataset is imbalanced, i.e., if one class has significantly more samples than the other class(es).

To perform stratified sampling, we can set the stratify parameter of the train_test_split function to the target variable:

1
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)

4. Wrapping it up

In this tutorial, we learned how to split a machine learning dataset into training and test sets in Python using the scikit-learn library. We saw how to use the train_test_split function to shuffle and split the dataset, and how to perform stratified sampling to ensure that the training and test sets have a similar class distribution to the original dataset. Splitting the dataset into training and test sets is an important step in the model building process, as it allows us to evaluate the model’s performance on unseen data.