# How to Split Data into Training and Test Sets Using Python

In machine learning, splitting your dataset into training and test sets is a critical step. It ensures that your model is trained on one part of the data and evaluated on a separate, unseen portion, enabling you to assess its real-world performance. In this article, we’ll walk through a Python script to efficiently split your data and explain its importance.

# **Why Split Data?**

When building machine learning models, it’s essential to avoid overfitting — a situation where the model performs well on training data but poorly on unseen data. By splitting the dataset, we achieve:

1. **Training Set:** Used to train the model.
    
2. **Test Set:** Used to evaluate the model’s performance.
    

# **Python Libraries for Data Splitting**

Python offers several tools for [data](https://nareshit.com/courses/data-science-online-training) splitting, but the most widely used is `train_test_split` from Scikit-learn’s `model_selection` module.

## **Key Features of** `train_test_split`

* **Random Sampling:** Ensures unbiased splitting.
    
* **Customizable Ratios:** You can specify the proportion of training and test data.
    
* **Stratification:** Ensures balanced class distribution in both sets for classification tasks.
    

# **Step-by-Step Script**

Below is a Python script to split data into training and test sets:

\# Import required libraries  
import pandas as pd  
from sklearn.model\_selection import train\_test\_split

\# Sample dataset (replace this with your actual dataset)  
data = {  
‘Feature1’: \[10, 20, 30, 40, 50, 60, 70, 80, 90, 100\],  
‘Feature2’: \[1, 2, 3, 4, 5, 6, 7, 8, 9, 10\],  
‘Label’: \[0, 1, 0, 1, 0, 1, 0, 1, 0, 1\]  
}

\# Convert the dictionary to a DataFrame  
df = pd.DataFrame(data)

\# Separate features (X) and target (y)  
X = df\[\[‘Feature1’, ‘Feature2’\]\]  
y = df\[‘Label’\]

\# Split the data into training and test sets  
X\_train, X\_test, y\_train, y\_test = train\_test\_split(  
X, y, test\_size=0.2, random\_state=42, stratify=y  
)

\# Print the sizes of each set  
print(“Training Features Shape:”, X\_train.shape)  
print(“Test Features Shape:”, X\_test.shape)  
print(“Training Labels Shape:”, y\_train.shape)  
print(“Test Labels Shape:”, y\_test.shape)

# **Explanation of the Code**

1. **Import Libraries:**
    

* `pandas` for data manipulation.
    
* `train_test_split` from Scikit-learn to perform the splitting.
    

**2\. Dataset Preparation:**

* Create a simple dataset (you can load your own dataset using `pd.read_csv()`).
    
* Separate the independent variables (`X`) and the dependent variable (`y`).
    

**3\. Splitting the Data:**

* `test_size=0.2`: Allocates 20% of the data to the test set and 80% to the training set.
    
* `random_state=42`: Ensures reproducibility by fixing the random seed.
    
* `stratify=y`: Maintains the class distribution of the target variable in both sets.
    

**4\. Check the Resulting Shapes:**

* The shapes of `X_train`, `X_test`, `y_train`, and `y_test` confirm the split.
    

![https://nareshit.com/courses/data-science-online-training](https://miro.medium.com/v2/resize:fit:700/1*AORZEB7Ig_0O71MTFOYOGA.png align="left")

How to Split Data into Training and Test Sets Using Python

# **Best Practices for Data Splitting**

1. **Set a Random State:** Always set a [`random_state`](https://nareshit.com/courses/data-science-online-training) for reproducibility, especially when sharing your work with others.
    
2. **Stratify for Classification Tasks:** Use `stratify=y` to ensure the class distribution in the training and test sets mirrors the original dataset.
    
3. **Consider Validation Sets:** For larger datasets, consider splitting further into training, validation, and test sets. The validation set helps fine-tune the model without overfitting to the test data.
    

# **Output Example**

If you run the script above, you’ll see output like this:

Training Features Shape: (8, 2)  
Test Features Shape: (2, 2)  
Training Labels Shape: (8,)  
Test Labels Shape: (2,)

This indicates that 8 rows (80%) are in the training set, and 2 rows (20%) are in the test set.

# **Conclusion**

Splitting data into training and test sets is a foundational step in building reliable machine learning models. By using Scikit-learn’s `train_test_split`, you can ensure that your data is divided effectively and reproducibly. With this knowledge, you’re one step closer to developing robust machine learning systems!

Let us know in the comments if you have any questions or need help with advanced splitting techniques.

For More Details Visit : [https://nareshit.com/courses/data-science-online-training](https://nareshit.com/courses/data-science-online-training)

Register For Free Demo on UpComing Batches : [https://nareshit.com/new-batches](https://nareshit.com/new-batches)