How to split data into train and test in Machine Learning

Splitting your data into train and test sets is one of the first steps in any Machine Learning project: you train the model on one set and evaluate it on another it has never seen, avoiding overly optimistic results. Below you will see how to split data into train and test in Machine Learning with scikit-learn's train_test_split function, using simple examples you can adapt to your own case.

Prerequisites

• Python 3 installed (version 3.9 or later).
• The scikit-learn and pandas libraries installed (pip install scikit-learn pandas).
• A dataset with input columns (features) and a target column.
• Basic Python knowledge, such as variables and importing modules.

Step 1: Prepare the data

Before splitting, we separate the features (the columns the model uses to predict, stored in X) from the target column (what we want to predict, stored in y). Let's use a small example with pandas.

import pandas as pd

# Example dataset
data = pd.DataFrame({
    "study_hours": [1, 2, 3, 4, 5, 6, 7, 8],
    "absences":    [8, 6, 5, 4, 3, 2, 1, 0],
    "passed":      [0, 0, 0, 1, 1, 1, 1, 1]
})

X = data[["study_hours", "absences"]]  # features
y = data["passed"]                      # target

By convention, X is a table with several columns and y is a single column. Keeping this separation clear prevents the target from accidentally leaking into the features — a mistake that would let the model "peek" at the answer.

Step 2: Split with train_test_split

The train_test_split function from the sklearn.model_selection module shuffles the data and separates it into two sets. You just pass X and y and set the percentage for testing.

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,     # 20% for testing, 80% for training
    random_state=42    # makes the split reproducible
)

Note the order of the variables on the left: first the two X sets (train and test) and only then the two y sets. Swapping this order is a common mistake that causes strange results later on.

Step 3: Understand test_size and random_state

Two parameters control how the split behaves:

• test_size: the fraction of data reserved for testing. 0.2 means 20%; alternatively you can set train_size.
• random_state: fixes the random seed. Using the same number every time (for example 42), you get exactly the same split each run, making your experiments reproducible.

By default, the function shuffles the rows before splitting (shuffle=True). This matters because many datasets come sorted — for example, with all positive cases at the end — and splitting without shuffling would leave the test set biased.

Step 4: Keep proportions with stratify

In classification problems, the train and test sets should have the same class proportions. If 30% of the cases belong to the "passed" class, we want that 30% in both sets. The stratify parameter ensures this:

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42,
    stratify=y         # keep class proportions of y
)

Tip: use stratify=y whenever the target variable is categorical. In regression problems (a continuous target), leave stratify out.

Check the result

To confirm the split worked, check the number of rows in each set with the .shape attribute and the class proportions:

print("Train:", X_train.shape, "Test:", X_test.shape)
print(y_train.value_counts(normalize=True))

The sum of train and test rows should equal the original total, and the class proportions should be similar in both sets. Because we fixed random_state, these numbers repeat every time you run the script.

Conclusion

In just a few lines you split the data into train and test — the step that guarantees an honest evaluation of a Machine Learning model. From here you can train an algorithm with X_train and y_train and measure performance with X_test and y_test, for example with a confusion matrix. As a next step, try changing test_size to 0.3 and watch how the set sizes change. What will be the ideal test percentage for your problem?