Data Splitting In Machine Learning

|

In machine learning, building an accurate predictive model is only half the battle — evaluating whether the model generalizes well to unseen data is equally critical. This is where data splitting comes into play. Data splitting involves dividing a dataset into separate subsets for training, validation, and testing to ensure that a model is evaluated fairly and avoids overfitting.

Without proper data splitting, a model may appear highly accurate on the data it has seen but fail miserably when faced with real-world examples. The practice ensures that we can measure true generalization performance instead of an overly optimistic estimate.

Why Data Splitting is Important

Machine learning models learn patterns from data. However, if we evaluate the model only on the same data it was trained on, we run the risk of memorization instead of generalization.

  • Overfitting: The model performs well on training data but poorly on new data.

  • Underfitting: The model performs poorly on both training and new data.

  • Proper generalization: The model learns robust patterns applicable to new situations.

By splitting data appropriately, we simulate how the model will behave in production. This is crucial for avoiding misleading performance metrics.

Common Types of Data Splits

There are three primary subsets commonly used in machine learning:

  1. Training Set:
    Used to train the machine learning algorithm. The model adjusts its parameters based on these data samples.

  2. Validation Set:
    Used to tune hyperparameters and make decisions during model development without touching the test data. Prevents data leakage from the final evaluation stage.

  3. Test Set:
    Used only once — after all training and hyperparameter tuning are complete — to provide an unbiased estimate of the model’s performance.

Basic Data Splitting with train_test_split

The simplest and most widely used function for data splitting in Python is train_test_split from scikit-learn. It allows us to divide the dataset into training and testing sets quickly.

Example:

from sklearn.model_selection import train_test_split
import numpy as np# Example dataset
X = np.arange(10).reshape((5, 2)) # Features
y = np.array([0, 1, 0, 1, 0]) # Labels# Split data: 70% train, 30% test
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.3, random_state=42
)print(“Training features:\n”, X_train)
print(“Test features:\n”, X_test)
print(“Training labels:”, y_train)
print(“Test labels:”, y_test)

  • test_size=0.3 → 30% of the data is reserved for testing.

  • random_state=42 → Ensures reproducible splits.

Adding a Validation Set

For more complex models, it’s often necessary to introduce a validation set. While train_test_split does not directly create three sets, we can achieve this with two consecutive calls.

Example:

from sklearn.model_selection import train_test_split

# First split: train+valid and test
X_temp, X_test, y_temp, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)

# Second split: train and validation
X_train, X_val, y_train, y_val = train_test_split(
X_temp, y_temp, test_size=0.25, random_state=42
)
# 0.25 x 0.8 = 0.2 → 20% validation

print(“Train size:”, len(X_train))
print(“Validation size:”, len(X_val))
print(“Test size:”, len(X_test))

Here:

  • Training: 60% of original data

  • Validation: 20% of original data

  • Test: 20% of original data

Stratified Splitting for Imbalanced Datasets

In classification tasks where classes are imbalanced (e.g., 90% class A, 10% class B), random splitting can lead to poor representation of minority classes. Stratified splitting ensures the same proportion of classes in all subsets.

Example:

from sklearn.model_selection import train_test_split

# Simulated imbalanced dataset
y = [0, 0, 0, 0, 1, 0, 0, 1, 0, 0]

X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.3, stratify=y, random_state=42
)

print(“Training label distribution:”, np.bincount(y_train))
print(“Test label distribution:”, np.bincount(y_test))

  • stratify=y → Preserves class distribution across training and test sets.

Cross-Validation: A Robust Alternative

While a single train/test split is simple, it may give an evaluation that depends heavily on the specific split. Cross-validation (CV) mitigates this by averaging performance across multiple folds.

  • K-Fold Cross-Validation: Data is split into k folds. The model is trained on k-1 folds and validated on the remaining fold. This process repeats k times, and results are averaged.

Example:

from sklearn.model_selection import KFold
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_scoreX, y = load_iris(return_X_y=True)
kf = KFold(n_splits=5, shuffle=True, random_state=42)
model = LogisticRegression(max_iter=200)accuracies = []
for train_index, test_index in kf.split(X):
X_train, X_test = X[train_index], X[test_index]
y_train, y_test = y[train_index], y[test_index]model.fit(X_train, y_train)
preds = model.predict(X_test)
accuracies.append(accuracy_score(y_test, preds))print(“Cross-validation accuracies:”, accuracies)
print(“Average accuracy:”, sum(accuracies) / len(accuracies))

Time Series Splitting

If your data has a temporal component (e.g., stock prices, weather data), random splitting may break the chronological order, causing data leakage. Instead, use time series splitting, which respects temporal order.

Example:

from sklearn.model_selection import TimeSeriesSplit
import numpy as npX = np.arange(10).reshape(-1, 1)
y = np.arange(10)
tscv = TimeSeriesSplit(n_splits=3)for train_index, test_index in tscv.split(X):
print(“TRAIN indices:”, train_index, “TEST indices:”, test_index)

Practical Guidelines for Data Splitting

  1. Always keep the test set separate and untouched until final evaluation.

  2. Use stratification for classification tasks to preserve label ratios.

  3. Use validation sets or cross-validation to tune hyperparameters.

  4. Maintain temporal order for time series data to avoid future leakage.

  5. Use a sufficiently large test set — typically 20–30% of data for small datasets.

Coding Example: Full Pipeline with Data Splitting

Here’s a complete example combining train-validation-test splitting, training a classifier, and evaluating its performance.

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score# Load dataset
data = load_breast_cancer()
X, y = data.data, data.target# Split train+val and test
X_temp, X_test, y_temp, y_test = train_test_split(
X, y, test_size=0.2, stratify=y, random_state=42
)# Split train and validation
X_train, X_val, y_train, y_val = train_test_split(
X_temp, y_temp, test_size=0.25, stratify=y_temp, random_state=42
)# Train model
model = RandomForestClassifier(random_state=42)
model.fit(X_train, y_train)# Validate
val_preds = model.predict(X_val)
print(“Validation accuracy:”, accuracy_score(y_val, val_preds))# Final test
test_preds = model.predict(X_test)
print(“Test accuracy:”, accuracy_score(y_test, test_preds))

This pipeline demonstrates:

  • Stratified splitting to handle class balance.

  • Separate validation set to tune the model before testing.

  • Final test accuracy as an unbiased performance estimate.

Conclusion

Data splitting is not just a preliminary step in machine learning — it is the foundation of reliable model evaluation. Whether you use a simple train/test division, a train-validation-test pipeline, or more advanced strategies like stratified sampling and cross-validation, the goal remains the same: ensuring that your model’s performance reflects how it will behave in the real world.

  • A training set teaches the model.

  • A validation set helps refine it without overfitting.

  • A test set holds the model accountable by measuring true generalization.

Failing to separate these properly can lead to overly optimistic results, wasted development time, and poor real-world performance. More sophisticated methods like cross-validation and time series splitting provide additional safeguards for specific use cases.

In practice, always respect data integrity:

  • For classification, stratify to maintain label distribution.

  • For time-dependent problems, never mix past and future data.

  • For hyperparameter tuning, avoid contaminating your test set.

By following these principles, you can confidently build models that are not just accurate on paper, but also robust and trustworthy in production.