Train Test Split

Code

import numpy as np
from sklearn.model_selection import train_test_split, StratifiedShuffleSplit

X = np.random.rand(100, 5)
y = np.random.randint(0, 2, 100)

# Standard split
X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, random_state=42, shuffle=True)

# Stratified split preserves class proportions
X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y)

# Three-way split: train/val/test
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42)

# Repeated stratified splits
sss = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=42)
for train_idx, test_idx in sss.split(X, y):
    print(len(train_idx), len(test_idx))

Explanation

train_test_split shuffles and divides data into training and evaluation subsets, with test_size controlling the proportion. Setting stratify=y keeps class balance, important for imbalanced classification. A three-way split nests two calls to obtain separate train, validation, and test sets.

More Machine Learning Snippets