Splitting a dataset

A brief explanation of how to do train-test split of a dataset using sklearn

Image by author

To train any machine learning model irrespective what type of dataset is being used you have to split the dataset into training data and testing data. So, let us look into how it can be done?

Here I am going to use the iris dataset and split it using the ‘train_test_split’ library from sklearn

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris

Then I load the iris dataset into a variable.

iris = load_iris()

Which I then use to store the data and target value into two separate variables.

x, y = iris.data, iris.target

Here I have used the ‘train_test_split’ to split the data in 80:20 ratio i.e. 80% of the data will be used for training the model while 20% will be used for testing the model that is built out of it.

x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.2,random_state=123)

As you can see here I have passed the following parameters in ‘train_test_split’:

1. x and y that we had previously defined
2. test_size: This is set 0.2 thus defining the test size will be 20% of the dataset
3. random_state: it controls the shuffling applied to the data before applying the split. Setting random_state a fixed value will guarantee that the same sequence of random numbers are generated each time you run the code.

When splitting a dataset there are two competing concerns:
-If you have less training data, your parameter estimates have greater variance.
-And if you have less testing data, your performance statistic will have greater variance.
The data should be divided in such a way that neither of them is too high, which is more dependent on the amount of data you have. If your data is too small then no split will give you satisfactory variance so you will have to do cross-validation but if your data is huge then it doesn’t really matter whether you choose an 80:20 split or a 90:10 split (indeed you may choose to use less training data as otherwise, it might be more computationally intensive).