How To Do Train Test Split Using Sklearn In Python
Last Updated :
27 Jun, 2022
In this article, let’s learn how to do a train test split using Sklearn in Python.
Train Test Split Using Sklearn
The train_test_split() method is used to split our data into train and test sets.
First, we need to divide our data into features (X) and labels (y). The dataframe gets divided into X_train,X_test , y_train and y_test. X_train and y_train sets are used for training and fitting the model. The X_test and y_test sets are used for testing the model if it’s predicting the right outputs/labels. we can explicitly test the size of the train and test sets. It is suggested to keep our train sets larger than the test sets.
Train set: The training dataset is a set of data that was utilized to fit the model. The dataset on which the model is trained. This data is seen and learned by the model.
Test set: The test dataset is a subset of the training dataset that is utilized to give an accurate evaluation of a final model fit.
validation set: A validation dataset is a sample of data from your model’s training set that is used to estimate model performance while tuning the model’s hyperparameters.
by default, 25% of our data is test set and 75% data goes into training tests.
Syntax: sklearn.model_selection.train_test_split()
parameters:
- *arrays: sequence of indexables. Lists, numpy arrays, scipy-sparse matrices, and pandas dataframes are all valid inputs.
- test_size: int or float, by default None. If float, it should be between 0.0 and 1.0 and represent the percentage of the dataset to test split. If int is used, it refers to the total number of test samples. If the value is None, the complement of the train size is used. It will be set to 0.25 if train size is also None.
- train_size: int or float, by default None.
- random_state : int,by default None. Controls how the data is shuffled before the split is implemented. For repeatable output across several function calls, pass an int.
- shuffle: boolean object , by default True. Whether or not the data should be shuffled before splitting. Stratify must be None if shuffle=False.
- stratify: array-like object , by default it is None. If None is selected, the data is stratified using these as class labels.
returns: splitting: list
Example 1:
The numpy, pandas, and scikit-learn packages are imported. The CSV file is imported. X contains the features and y is the labels. we split the dataframe into X and y and perform train test split on them. random_state acts like a numpy seed, it is used for data reproducibility. test_size is given as 0.25 , it means 25% of our data goes into our test size. 1-test_size is our train size, we don’t need to specify that. shuffle =True, shuffles our data before spilling. The X_train and X_test sets are used to fit and train our model and the test sets are used for testing and validating.
To access the CSV file click here.
Python3
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
df = pd.read_csv('headbrain1.csv')
print(df.head())
X= df['Head Size(cm^3)']
y=df['Brain Weight(grams)']
X_train, X_test,
y_train, y_test = train_test_split(X,y ,
random_state=104,
test_size=0.25,
shuffle=True)
print('X_train : ')
print(X_train.head())
print('')
print('X_test : ')
print(X_test.head())
print('')
print('y_train : ')
print(y_train.head())
print('')
print('y_test : ')
print(y_test.head())
|
Output:
Head Size(cm^3) Brain Weight(grams)
0 4512 1530
1 3738 1297
2 4261 1335
3 3777 1282
4 4177 1590
X_train :
99 3478
52 4270
184 3479
139 3171
107 3399
Name: Head Size(cm^3), dtype: int64
(177,)
X_test :
66 3415
113 3594
135 3436
227 4204
68 4430
Name: Head Size(cm^3), dtype: int64
(60,)
y_train :
99 1270
52 1335
184 1160
139 1127
107 1226
Name: Brain Weight(grams), dtype: int64
(177,)
y_test :
66 1310
113 1290
135 1235
227 1380
68 1510
Name: Brain Weight(grams), dtype: int64
(60,)
Example 2:
In this example, the same steps are followed, instead of specifying the test_size we specify the train_size. test_size is 1-train_size. 80% of the data is train set, so 20% of our data is our test set. If we don’t specify the sizes of test and train sets by default test_size will be 0.25. X_train and y_train have the same shape and indexes, as y_train is the label for X_train features. same goes with X_test and y_test.
Python3
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
df = pd.read_csv('headbrain1.csv')
print(df.shape)
print(df.head())
X= df['Head Size(cm^3)']
y=df['Brain Weight(grams)']
X_train, X_test, y_train,
y_test = train_test_split(X,y ,
random_state=104,
train_size=0.8, shuffle=True)
print('X_train : ')
print(X_train.head())
print(X_train.shape)
print('')
print('X_test : ')
print(X_test.head())
print(X_test.shape)
print('')
print('y_train : ')
print(y_train.head())
print(y_train.shape)
print('')
print('y_test : ')
print(y_test.head())
print(y_test.shape)
|
Output:
(237, 2)
Head Size(cm^3) Brain Weight(grams)
0 4512 1530
1 3738 1297
2 4261 1335
3 3777 1282
4 4177 1590
X_train :
110 3695
164 3497
58 3935
199 3297
182 4005
Name: Head Size(cm^3), dtype: int64
(189,)
X_test :
66 3415
113 3594
135 3436
227 4204
68 4430
Name: Head Size(cm^3), dtype: int64
(48,)
y_train :
110 1310
164 1280
58 1330
199 1220
182 1280
Name: Brain Weight(grams), dtype: int64
(189,)
y_test :
66 1310
113 1290
135 1235
227 1380
68 1510
Name: Brain Weight(grams), dtype: int64
(48,)
Like Article
Suggest improvement
Share your thoughts in the comments
Please Login to comment...