Skip to main content

Data Splitting Techniques

Data splitting is a crucial step in the data science pipeline, serving as the foundation for model evaluation and ensuring that your machine learning models generalize well to unseen data. This article explores the theory and practice of train/test/validation splitting, with a focus on practical implementation using pandas and NumPy.

1. Introduction​

1.1 What is Data Splitting?​

Data splitting refers to the process of dividing your dataset into distinct subsets that serve different purposes during the machine learning workflow. The most common splits are:

  • Training Set: Used to train the model.
  • Validation Set: Used to tune the model's hyperparameters and assess model performance during training.
  • Test Set: Used to evaluate the model's performance on unseen data.

1.2 Why is Data Splitting Important?​

Splitting your data helps to:

  • Prevent Overfitting: By evaluating model performance on separate data, you can avoid overfitting to the training data.
  • Ensure Generalization: Splits allow you to assess how well your model will perform on new, unseen data.
  • Optimize Hyperparameters: Validation sets enable fine-tuning of model parameters to improve performance.

2. Train/Test Split​

2.1 Basic Train/Test Split​

A basic train/test split involves dividing your dataset into two subsets: one for training and one for testing.

Example: Basic Train/Test Split​

import numpy as np
import pandas as pd

# Sample DataFrame
rng = np.random.default_rng(42)
df = pd.DataFrame({
    'Feature1': rng.standard_normal(100),
    'Feature2': rng.standard_normal(100),
    'Target': np.tile([0, 1], 50)
})

# Shuffle the dataset
df_shuffled = df.sample(frac=1, random_state=42).reset_index(drop=True)

# Split the dataset: 80% train, 20% test
train_size = int(0.8 * len(df_shuffled))
train_set = df_shuffled[:train_size]
test_set = df_shuffled[train_size:]

print("Train set size:", len(train_set))
print("Test set size:", len(test_set))

Output:

Train set size: 80
Test set size: 20

2.2 Considerations for Train/Test Split​

  • Shuffling: Shuffle independent and identically distributed observations when row order has no meaning. Do not randomly shuffle temporal observations or records that belong to the same person, device, household, or other group.
  • Stratification: For classification, stratification can stabilize class proportions when every class has enough examples. It does not replace group-aware or time-aware splitting and can make the apparent variability between folds artificially small.

3. Train/Validation/Test Split​

3.1 Adding a Validation Set​

A more robust approach involves splitting the data into three subsets: training, validation, and test. The validation set is used to fine-tune model parameters and avoid overfitting during the training phase.

Example: Train/Validation/Test Split​

# Further split the training set into training and validation sets: 70% train, 10% validation, 20% test
val_size = int(0.1 * len(df_shuffled))
train_set = df_shuffled[:train_size - val_size]
val_set = df_shuffled[train_size - val_size:train_size]
test_set = df_shuffled[train_size:]

print("Train set size:", len(train_set))
print("Validation set size:", len(val_set))
print("Test set size:", len(test_set))

Output:

Train set size: 70
Validation set size: 10
Test set size: 20

3.2 Importance of Validation Sets​

  • Hyperparameter Tuning: The validation set allows for fine-tuning of hyperparameters without contaminating the test set, which should only be used for final model evaluation.
  • Preventing Overfitting: By monitoring performance on the validation set, you can stop training once the model starts to overfit (e.g., when performance on the validation set begins to degrade).

4. Cross-Validation Techniques​

4.1 What is Cross-Validation?​

Cross-validation is a technique that involves dividing the dataset into multiple subsets (folds) and training the model multiple times, each time using a different subset as the validation set and the remaining data as the training set. This approach provides a more accurate estimate of model performance.

4.2 K-Fold Cross-Validation​

In K-Fold Cross-Validation, the dataset is split into K equal-sized folds. The model is trained K times, each time leaving out one fold as the validation set and training on the remaining K-1 folds.

Example: Implementing K-Fold Cross-Validation​

from sklearn.model_selection import StratifiedKFold

# Create five reproducible folds while preserving class proportions.
splitter = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

# Inspect the fold boundaries; an estimator would be fitted on each train_set.
for i, (train_index, val_index) in enumerate(
    splitter.split(df_shuffled, df_shuffled['Target'])
):
    train_set = df_shuffled.iloc[train_index]
    val_set = df_shuffled.iloc[val_index]
    
    # Print sizes of the training and validation sets
    print(f"Fold {i + 1}:")
    print("Train set size:", len(train_set))
    print("Validation set size:", len(val_set))
    print()

Output:

Fold 1:
Train set size: 80
Validation set size: 20

Fold 2:
Train set size: 80
Validation set size: 20

...

Fold 5:
Train set size: 80
Validation set size: 20

4.3 Leave-One-Out Cross-Validation (LOOCV)​

LOOCV is an extreme form of cross-validation where each fold contains a single data point. The model is trained on all other data points and tested on the single point left out.

Example: Implementing LOOCV​

# LOOCV manually (for demonstration)
for i in range(len(df_shuffled)):
    # Leave one out for validation
    val_set = df_shuffled.iloc[[i]]
    train_set = df_shuffled.drop(i)
    
    # Print sizes of the training and validation sets
    print(f"Iteration {i + 1}:")
    print("Train set size:", len(train_set))
    print("Validation set size:", len(val_set))
    print()

Output:

Iteration 1:
Train set size: 99
Validation set size: 1

Iteration 2:
Train set size: 99
Validation set size: 1

...

Iteration 100:
Train set size: 99
Validation set size: 1

4.4 Considerations for Cross-Validation​

  • Computational Cost: Cross-validation, particularly LOOCV, can be computationally expensive.
  • Bias-Variance Tradeoff: K-Fold Cross-Validation strikes a good balance between bias and variance, with typical values of K being 5 or 10.

5. Best Practices for Data Splitting​

5.1 Avoiding Data Leakage​

Data leakage occurs when information unavailable at prediction time influences training or model selection, leading to overly optimistic estimates. Split first. Then learn imputation values, scaling parameters, encodings, feature selection, dimensionality reduction, and any resampling only from the training partition. During cross-validation, put those learned steps in a pipeline so each step is fitted again inside each training fold. Keep the test set isolated until one final evaluation; repeatedly checking it turns it into validation data.

5.2 Ensuring Representativeness​

  • Stratification: Consider stratification for classification when class proportions matter and it is compatible with the sampling unit.
  • Groups: Keep related observations in the same partition with a group-aware splitter; otherwise the model can effectively see the same entity during training and evaluation.
  • Randomization: Shuffle only exchangeable observations. Preserve order for forecasting and respect the real deployment boundary.

5.3 Split Proportions​

Common split proportions include:

  • Train/Test: 80/20 or 70/30
  • Train/Validation/Test: 70/15/15 or 60/20/20
  • K-Fold: Typically 5 or 10 folds

5.4 Handling Time Series Data​

For time series data, traditional random splitting is inappropriate because it breaks the temporal order. Instead:

  • Train/Test Split: Use the earlier data for training and the later data for testing.
  • Rolling Cross-Validation: Train on a rolling window of data, increasing the size of the training set with each iteration.

6. Conclusion​

6.1 Summary of Key Techniques​

In this article, we've covered the theory and practice of data splitting, including basic train/test splits, the addition of validation sets, and cross-validation techniques. Properly splitting your data is essential to building robust machine learning models that generalize well to unseen data.

6.2 Further Reading​

By mastering data splitting techniques, you ensure that your models are evaluated rigorously, reducing the risk of overfitting and improving their generalization to real-world data.