Skip to main content
Taught by Tech Leads

Master pipelines, cloud & AI to become an operational Data Engineer.

DataScientist.fr
Image de Split your dataset using train_test_split from scikit-learn - Practical tutorial in Python
Data Scientist
Artificial Intelligence
Python

Split your dataset using train_test_split from scikit-learn - Practical tutorial in Python

Photo de Romain DE LA SOUCHÈRE

Tech Lead, CTO AXI Technologies

Published on 2 janvier 2025 · 12 min of reading

In the world of machine learning, the key to a high-performing model often lies in how the data is split. A judicious separation between training, validation, and test sets is crucial to avoid the pitfalls of underfitting and overfitting. This article explores the importance of these splits and presents effective methods to optimize the use of train_test_split, thereby ensuring robust and reliable results in supervised learning. Dive into the subtleties of this fundamental step to strengthen your predictive models.

The Importance of Data Splitting

To understand the importance of data splitting, it is essential to grasp how this step influences the performance and reliability of machine learning models. Data splitting is a common practice that involves separating a dataset into two subsets: a training set and a test set. This separation is crucial for assessing a model's ability to generalize to unseen data.

Preventing Overfitting

One of the main reasons to split the data is to prevent overfitting. When you train a model on the entirety of your dataset, it risks fitting too closely to the peculiarities of that set, capturing noise instead of grasping general trends. By splitting the data, you can assess the model's performance on a test set that was not used during training, thus ensuring that the model generalizes well.

Performance Evaluation

Data splitting also allows for an objective evaluation of a model's performance. Without a separate test set, it would be difficult to determine if a model performs well simply because it fits the training data or if it can truly make accurate predictions on new data. By using metrics such as accuracy, precision, recall, or the F1 score on the test set, you obtain a more reliable indication of your model's predictive capabilities.

Hyperparameter Tuning

Another key reason for data splitting is hyperparameter tuning. Hyperparameters are parameters set before the model training process, and their tuning can significantly impact model performance. By using a technique like cross-validation, where the dataset is divided into several parts, you can test different model configurations and choose the one that offers the best performance on the validation set without contaminating the test set.

Balanced Distribution

Finally, it is important to ensure that data splitting is done in a way that maintains a balanced distribution of classes in both the training and test sets. This guarantees that the model is trained and evaluated fairly, especially in cases of imbalanced data where certain classes may be underrepresented. An imbalance could skew performance evaluation, leading to erroneous conclusions.
In conclusion, data splitting is a fundamental step in developing robust and reliable machine learning models.

Training, Validation, and Test Sets

When working with machine learning models, it is essential to understand the specific roles of training, validation, and test sets. Each of these sets has a distinct purpose that contributes to the creation of effective and reliable models.

Training Set

The training set is the portion of the data used to adjust the model's internal parameters. It is on this set that the model 'learns' from examples, minimizing a cost function that measures the error between the model's predictions and the actual values. To maximize learning, it is often desirable for this set to be as large as possible while retaining a sufficient portion for validation and test sets.

Validation Set

The validation set plays a crucial role in tuning the model's hyperparameters and preventing overfitting. By evaluating the model on this set during training, you can adjust aspects such as the learning rate, the number of layers in a neural network, or the regularization parameters. Cross-validation, often employed, involves dividing the dataset into several subsets and rotating the training and validation roles to robustly optimize these hyperparameters without using the test set.

Test Set

The test set is used only at the end of the modeling process to definitively assess the model's performance. It provides an unbiased evaluation of the model's generalization ability on unseen data. It is crucial not to expose this set during the training and validation process to avoid biasing the results and overestimating the model's performance.

Best Practices

To ensure accurate evaluation, a good practice is to use a typical split of 70% for training, 15% for validation, and 15% for testing. However, this distribution may vary depending on the size of the dataset and the complexity of the model. It is also recommended to ensure that each set is representative of the overall dataset, especially in the case of imbalanced data.
In summary, distinguishing between these sets and their appropriate use is fundamental to building robust and reliable models.

Underfitting and Overfitting

In the field of machine learning, underfitting and overfitting are two common problems that can affect a model's performance. Understanding these concepts is crucial for effectively tuning models.

Underfitting

Underfitting occurs when the model is too simple to capture the underlying trends in the data. This can result from using a model with too few parameters or insufficient complexity. For example, attempting to model a complex nonlinear relationship with linear regression can lead to underfitting.
Indicators of underfitting include poor performance on both training and validation sets. To remedy this, it is often necessary to increase the model's complexity, either by adding more parameters or opting for a more sophisticated model, such as a deeper neural network or an ensemble model.

Overfitting

Overfitting occurs when the model fits the training data too well, capturing noise and minor variations that do not generalize well to new data. This typically results in excellent performance on the training set but poor performance on the validation and test sets.
To detect overfitting, it is helpful to monitor the performance gap between the training and validation sets. If the model exhibits significantly better performance on the training set than on the validation set, overfitting is likely occurring.

Mitigation Strategies

Several strategies can be implemented to mitigate overfitting:
  • Regularization: Adding regularization terms to the cost function to penalize complex models.
  • Pruning: In decision trees, the tree size can be reduced to prevent it from fitting too closely to the training data.
  • Data Augmentation: Generating more training data through transformations to enhance the model's generalization capability.
By carefully managing underfitting and overfitting, you can develop models that not only fit well to your training data but also generalize effectively to new data.

Prerequisites for Using train_test_split

To effectively use the train_test_split function from the Scikit-learn library in Python, certain prerequisites and considerations must be taken into account. This function is essential for splitting your dataset into training and test subsets, and sometimes validation, to reliably assess your model's performance.

Understanding the Data

Before splitting your data, it is crucial to thoroughly understand the structure and characteristics of your dataset. This includes checking data types, handling missing values, and understanding class distribution in classification tasks. A good understanding will help you decide on the appropriate proportion for training and test sets.

Data Preprocessing

Before applying train_test_split, data preprocessing may be necessary. This preprocessing can include normalizing or standardizing features, converting categories into numerical variables, or handling outliers. These steps ensure that the data is ready for effective learning and that the training and test sets are comparable.

Parameters of train_test_split

Using train_test_split involves specifying certain parameters:
  • test_size: Determines the proportion of the dataset to include in the test set. For example, test_size=0.2 means that 20% of the data will be used for testing.
  • random_state: Sets the seed of the random number generator to obtain reproducible results. This is particularly useful for comparing the performance of different models consistently.
  • stratify: Ensures that the class distribution is similar in the training and test sets, which is crucial for imbalanced datasets. Use stratify=y where y is the class vector.

Executing the Function

Here is an example code snippet to illustrate the use of train_test_split:
python
By following these prerequisites and considerations, you can use train_test_split to efficiently divide your data and prepare your model for robust evaluations.

Application of train_test_split

The application of the train_test_split function from Scikit-learn is a key step in the data preparation process for machine learning. It allows for the random and controlled separation of a dataset into training and test subsets.

Practical Example

Imagine you are working with a classification dataset to predict whether a customer will purchase a product. This dataset contains features such as age, income, and purchase history.
python
In this example, we used test_size=0.25 to specify that 25% of the data should be reserved for testing. The random_state=42 parameter ensures the reproducibility of results.

Importance of Stratification

Stratification is particularly important in the case of imbalanced data. Suppose our y vector represents a binary classification problem with a strong predominance of class 0. By using stratification, we ensure that the proportions of each class are maintained in both training and test sets:
python

Verifying the Sets

After the split, it is wise to check the dimensions of the resulting sets to ensure that the separation was done correctly:
python
This quick verification ensures that your sets are correctly configured for the training and testing phase. By following these steps, you can ensure that train_test_split is applied optimally to prepare your data for machine learning.

Supervised Learning with train_test_split

Supervised learning is a paradigm of machine learning where a model is trained to make predictions based on labeled data. The use of train_test_split plays an essential role in this process by allowing for an effective assessment of the model's performance.

Model Training

With train_test_split, the training set is used to fit the model. For example, if you are using logistic regression for a binary classification task, you might train your model as follows:
python
In this scenario, the model learns the relationships in the training data to make predictions on new data.

Model Evaluation

Once the model is trained, the test set is used to evaluate its performance. This allows for measuring the model's generalization capability. For instance, you can calculate the model's accuracy with:
python
python
This step is crucial to ensure that the model does not suffer from overfitting and can generalize well to new data.

Hyperparameter Tuning

The use of train_test_split also facilitates hyperparameter tuning. By repeating the process of splitting, training, and testing with different parameter configurations, you can optimize model performance.
Using the subsets generated by train_test_split, supervised learning becomes a structured process that ensures the model is not only well-trained but also properly evaluated for effective deployment.

Conclusion

Throughout this article, we have explored in depth the importance and application of the train_test_split function in the context of supervised learning. This data splitting method is essential to ensure that machine learning models are not only well-trained but also capable of generalizing effectively to new data.

Summary of Key Points

We began by highlighting the importance of data splitting, which prevents overfitting and ensures a fair evaluation of model performance. Next, we distinguished the roles of training, validation, and test sets, each playing a crucial role in developing a robust model.
The analysis of underfitting and overfitting issues shed light on potential challenges in model training and strategies to prevent them, such as regularization and data augmentation.
We then detailed the prerequisites for using train_test_split, including understanding the data and necessary preprocessing to ensure that the split is effective and representative.

Application in Supervised Learning

The application of train_test_split in supervised learning was illustrated through the example of logistic regression, demonstrating how this function facilitates training and evaluating models. Through a judicious split of the data, models can be reliably assessed, allowing for performance optimization through hyperparameter tuning.

Importance of Validation

We also emphasized the importance of cross-validation as a complement to train_test_split, especially when fine-tuning hyperparameters without compromising the integrity of the final test set.
By integrating these practices into your workflow, you can ensure that your models are not only performing well on your training data but are also ready to tackle real-world data effectively. Through rigorous evaluation and well-structured training, you can develop machine learning solutions that provide real added value.

Want to go further?

This topic is part of our Become a Data Analyst course. Browse the full programme, or get it by email.

Share with

Photo de Romain DE LA SOUCHÈRE

Romain DE LA SOUCHÈRE

Tech Lead, CTO AXI Technologies

Expert Data Engineering et Cloud, Romain affiche plus de 11 ans d'expérience, dont plusieurs années comme Lead Developer sur des solutions Smart Building haute performance. Il y a conçu et mis en production des moteurs de traitement capables d'absorber des centaines de milliers de données de capteurs par minute, ainsi que des bases clusterisées gérant plus de 10 millions de données dynamiques. Certifié Microsoft Azure DevOps Engineer Expert, il maîtrise aussi bien le développement back-end (Python, C#) que le DevOps (Docker, Kubernetes, Terraform) et les agents LLM. Formateur en Python, cloud, DevOps et IA générative appliquée, il forme avec une obsession : Amener chaque apprenant à concevoir et déployer des architectures réellement scalables en production.

» Learn More

Associated trainings

All our trainings
Image de la formation Become a Data Analyst
Become a Data Analyst
6 months
Intermediate
Guarantee
Image de la formation Become a Data Engineer
Become a Data Engineer
9 months
Advanced
Guarantee