Train and Test Set in Python Machine Learning – How to Split

In our last session, we discussed Data Preprocessing, Analysis & Visualization in Python ML. Now, in this tutorial, we will learn how to split a CSV file into Train and Test Data in Python Machine Learning. Moreover, we will learn the prerequisites and process for splitting a dataset into training data and a test set in Python ML.

We’ve cleaned our data and seen the charts; now it’s time to focus on the ultimate stress test. Splitting your dataset is like an insurance policy for your code; it’s the only way you can prove that a machine learning model actually starts working before you make it work.

So, let’s begin with how to Train & Test Set in Python Machine Learning.

Training and Test Data in Python Machine Learning

As we work with datasets, a machine learning algorithm works in two stages. We usually split the data around 20%-80% between testing and training stages. Under supervised learning, we split a dataset into training data and test data in Python ML.

a. Prerequisites for Train and Test Data

We will need the following Python libraries for this tutorial: pandas and sklearn.
We can install these with pip-

pip install pandas

pip install sklearn

Also, We use pandas to import the dataset and sklearn to perform the splitting. You can import these packages as-

>>> import pandas as pd
>>> from sklearn.model_selection import train_test_split
>>> from sklearn.datasets import load_iris

Do you Know about Python Data File Formats – How to Read CSV, JSON, XLS 

How to Split Train and Test Set in Python Machine Learning?

The following is the process of Train and Test set in Python ML. So, let’s take a dataset first.

a. Loading the Dataset

Let’s load the forestfires dataset using pandas.

>>> data=pd.read_csv('forestfires.csv')
>>> data.head()

b. Splitting

Let’s split this data into labels and features. Now, what’s that? Using features, we predict labels. I mean using features (the data we use to predict labels), we predict labels (the data we want to predict).

>>> y=data.temp
>>> x=data.drop('temp',axis=1)

Temp is a label to predict temperatures in y; we use the drop() function to take all other data in x. Then, we split the data.

>>> x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.2)
>>> x_train.head()

>>> x_train.shape

(413, 12)
Do you Know How to work with relational databases with Python

>>> x_test.head()

>>> x_test.shape

(104, 12)
The line test_size=0.2 suggests that the test data should be 20% of the dataset and the rest should be train data. With the outputs of the shape() functions, you can see that we have 104 rows in the test data and 413 in the training data.

c. Another Example

Let’s take another example. We’ll use the IRIS dataset this time.

>>> iris=load_iris()
>>> x,y=iris.data,iris.target
>>> x_train,x_test,y_train,y_test=train_test_split(x,y,
train_size=0.5,
test_size=0.5,
random_state=123)
>>> y_test

array([1, 2, 2, 1, 0, 2, 1, 0, 0, 1, 2, 0, 1, 2, 2, 2, 0, 0, 1, 0, 0, 2,
0, 2, 0, 0, 0, 2, 2, 0, 2, 2, 0, 0, 1, 1, 2, 0, 0, 1, 1, 0, 2, 2,
2, 2, 2, 1, 0, 0, 2, 0, 0, 1, 1, 1, 1, 2, 1, 2, 0, 2, 1, 0, 0, 2,
1, 2, 2, 0, 1, 1, 2, 0, 2])

>>> y_train

array([1, 1, 0, 2, 2, 0, 0, 1, 1, 2, 0, 0, 1, 0, 1, 2, 0, 2, 0, 0, 1, 0,
0, 1, 2, 1, 1, 1, 0, 0, 1, 2, 0, 0, 1, 1, 1, 2, 1, 1, 1, 2, 0, 0,
1, 2, 2, 2, 2, 0, 1, 0, 1, 1, 0, 1, 2, 1, 2, 2, 0, 1, 0, 2, 2, 1,
1, 2, 2, 1, 0, 1, 1, 2, 2])

One of the essential approaches when dividing data into training and testing sets is to capture the model’s performance on unseen data. This practice aids in evaluating the model’s performance and potentially recognizing problems such as overfitting or underfitting. We continuously split the dataset in order to check our model’s capabilities of not memorizing the data, yet capable of classifying new instances.

Furthermore, it should be noted that setting the random_state parameter in train_test_split makes the dataset creation repeatable. This parameter is set to a fixed integer to make sure that the same splitting of data is achieved every time a specific model is tested; thus is essential during the testing of the model.

Let’s explore Python Machine Learning Environment Setup 

Plotting of Train and Test Set in Python

We fit our model on the training data to make predictions on it. Let’s import the linear_model from sklearn, apply linear regression to the dataset, and plot the results.

>>> from sklearn.linear_model import LinearRegression as lm
>>> model=lm().fit(x_train,y_train)
>>> predictions=model.predict(x_test)
>>> import matplotlib.pyplot as plt
>>> plt.scatter(y_test,predictions)

<matplotlib.collections.PathCollection object at 0x0651CA30>

>>> plt.xlabel('True values')

Text(0.5,0,’True values’)

>>> plt.ylabel('Predictions')

Text(0,0.5,’Predictions’)
Read about Python NumPy – NumPy ndarray & NumPy Array

>>> plt.show()

0.9396299518034936
So, this was all about Train and Test Set in Python Machine Learning. Hope you like our explanation.

Conclusion

Today, we learned how to split a CSV or a dataset into two subsets- the training set and the test set in Python Machine Learning. We usually let the test set be 20% of the entire data set, and the rest 80% will be the training set. Furthermore, if you have a query, feel free to ask in the comment box.
Related Topic- Python Geographic Maps & Graph Data
For reference