How to Split the dataset with scikit-learn’s train_test_split() in Python

Dataset splitting plays a crucial role in machine learning. It helps us to evaluate the performance of the model. In this tutorial, we will learn how to split the dataset using scikit-learn.

Splitting the dataset using scikit-learn

Steps involved:

  • Importing packages
  • Loading the dataset
  • Splitting using sklearn

Importing the packages:

import pandas as pd
from sklearn.model_selection import train_test_split

For splitting we need to import train_test_split from sklearn.

Loading the dataset:

Lets consider Sample.csv as the dataset

df = pd.read_csv("PATH OF THE DATASET")
(614, 13)
Index(['Loan_ID', 'Gender', 'Married', 'Dependents', 'Education',
'Self_Employed', 'ApplicantIncome', 'CoapplicantIncome', 'LoanAmount',
'Loan_Amount_Term', 'Credit_History', 'Property_Area', 'Loan_Status'],

In the dataset we can find that Loan_Status is dependent variable.

X = df.drop(['Loan_Status'],1)
(614, 12)
y = df['Loan_Status']

User input:

print("Enter the splitting factor:")
n = float(input())
Enter the splitting factor: 

Here user needs to give the factor by which train data and test data should be splitted. Let us consider 0.3 as splitting factor.

Splitting using sklearn:

from sklearn.model_selection import train_test_split
x_train, x_test, y_train, y_test = train_test_split(X,y, test_size=n)

Here we are splitting the dataset randomly into x_train, x_test, y_train, and y_test by given splitting factor.
NOTE: train_test_split(X,y, test_size=n, random_state = any integer) produces same result after every execution. Where as train_test_split(X,y, test_size=n) produces different results for every execution.

Before Splitting:

print("Size of x:")
print("Size of y:")
Size of x:
(614, 12)
Size of y:

After Splitting:

print("Size of x_train:")
print("Size of y_train:")
print("Size of x_test:")
print("Size of y_test:")
Size of x_train
(429, 12)
Size of y_train
Size of x_test
(185, 12)
Size of y_test

As the splitting factor is 0.3, 30% of total dataset ((i.e) 30% of 614 = 185) goes to test data and remaining goes to train successfully.
In this way the dataset is splitted into train and test using scikit-learn.

Also read:

Leave a Reply

Your email address will not be published. Required fields are marked *