Train-Test Split

1. Introduction

This final topic of Module 3 puts into practice a concept you were introduced to conceptually back in Module 2 (Training, Validation & Testing) — now with a hands-on focus on actually implementing the split correctly in code, including a few practical details (like stratification) that matter greatly in real projects.

2. What is Train-Test Split?

Simple definition: Train-test split is the practical step of dividing your dataset into a training portion (used to teach the model) and a testing portion (used to check how well it performs on new, unseen data).

Technical explanation: train_test_split(), provided by Scikit-learn, randomly partitions a dataset's features (X) and labels (y) into training and testing subsets according to a specified proportion, optionally preserving the original class distribution (via stratification) and ensuring reproducibility through a fixed random seed.

3. Why is it Important?

• It's the final, practical preprocessing step before actually training a model — everything in this module builds toward this point.
• Getting this step right (proper randomization, appropriate stratification) directly affects how trustworthy your model evaluation will be.
• It's used in virtually every single ML project you'll ever build.

4. Prerequisites

Comfort with Training, Validation & Testing (Module 2, Topic 8) and Features & Labels (Module 2, Topic 9).

5. Core Concepts

1. train_test_split() from Scikit-learn
2. The test_size parameter
3. The random_state parameter (reproducibility)
4. Stratified splitting for classification problems

6. Detailed Explanation

a) `train_test_split()`

This function takes your features (X) and labels (y) and returns four separate arrays: X_train, X_test, y_train, y_test — ready to be used directly with any Scikit-learn model.

b) The `test_size` Parameter

test_size specifies what proportion of the data should go into the test set (e.g., test_size=0.2 reserves 20% of the data for testing, leaving 80% for training).

c) The `random_state` Parameter

Setting a fixed random_state ensures that the split is reproducible — running the same code again produces the exact same train/test split, which is important for consistent, comparable experiments.

d) Stratified Splitting

For classification problems, especially with imbalanced classes (e.g., 90% "not fraud," 10% "fraud"), a plain random split might accidentally produce a test set with very few (or zero) examples of the minority class. Setting stratify=y ensures the train and test sets preserve the same class proportions as the original dataset.

7. How It Works

1. Prepare your fully cleaned, encoded, and scaled features (X) and label (y).
2. Call train_test_split(X, y, test_size=..., random_state=..., stratify=...).
3. Receive four outputs: X_train, X_test, y_train, y_test.
4. Use X_train/y_train to train your model, and reserve X_test/y_test for final evaluation only.

8. Real-World Example

Imagine a dataset for detecting rare diseases, where only 2% of patients actually have the disease. A plain random 80/20 split might, by chance, place very few (or even zero) actual disease cases into the test set, making it impossible to properly evaluate how well the model detects the disease. Using stratify=y ensures both the training and test sets maintain that same 2% disease rate, leading to a much more reliable evaluation.

9. Python Example

python
from sklearn.model_selection import train_test_split
import numpy as np

# Features and labels (imbalanced example: mostly 0s, few 1s)
X = np.array([[i] for i in range(20)])
y = np.array([0]*17 + [1]*3)   # 17 zeros, 3 ones (imbalanced)

# Plain random split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=42
)
print("Without stratification:")
print("y_train distribution:", np.bincount(y_train))
print("y_test distribution:", np.bincount(y_test))

# Stratified split
X_train_s, X_test_s, y_train_s, y_test_s = train_test_split(
    X, y, test_size=0.25, random_state=42, stratify=y
)
print("\nWith stratification:")
print("y_train distribution:", np.bincount(y_train_s))
print("y_test distribution:", np.bincount(y_test_s))

Expected Output (exact numbers may vary slightly, but stratified proportions will be preserved):

text
Without stratification:
y_train distribution: [13  2]
y_test distribution: [4 1]

With stratification:
y_train distribution: [13  2]
y_test distribution: [4 1]

(Note: with very small datasets like this example, both approaches can sometimes look similar by chance — the benefit of stratification becomes much clearer and more consistent with larger, more imbalanced real-world datasets.)

10. Code Explanation

• train_test_split(X, y, test_size=0.25, random_state=42) splits the data, reserving 25% for testing.
• np.bincount(y_train) counts how many of each class (0 and 1) ended up in the training set, letting us check whether the class balance was preserved.
• Adding stratify=y explicitly tells train_test_split() to maintain the same class proportions in both the training and test sets as in the original full dataset — this becomes especially important with larger, more imbalanced datasets than this small example.
• random_state=42 ensures this exact split can be reproduced consistently across multiple runs.

11. Advantages

• Provides a simple, one-line way to create reliable training and test sets.
• Stratification helps ensure fair, representative splits for classification problems with imbalanced classes.
• random_state ensures reproducibility, critical for debugging and fair experiment comparison.

12. Limitations

• A single random split (even with stratification) can still occasionally be somewhat unrepresentative, especially with very small datasets — Cross-Validation (Module 6) helps address this more robustly.
• Doesn't account for time-based dependencies in data (e.g., for time-series problems, a random split can leak future information into training — special time-based splitting is needed instead).

13. Common Mistakes

• Forgetting to use stratify=y for imbalanced classification problems, risking an unrepresentative test set.
• Not setting a random_state, making experiments inconsistent and hard to reproduce or compare.
• Applying train_test_split() to already-scaled data without considering that scaling should ideally happen after splitting (fit scaler on training data only, per Topic 7).
• Using a plain random split for time-series data, which can improperly mix future and past information.

14. Best Practices

• Always set a random_state for reproducible experiments.
• Use stratify=y for classification problems, especially with imbalanced classes.
• Perform the train-test split BEFORE fitting scalers or imputers, to avoid data leakage (fit only on the resulting training set).
• For time-series data, use a chronological split instead of a random one.

15. Real-World Applications

• The standard final preprocessing step before training virtually any supervised ML model.
• Essential for fair evaluation in fraud detection, medical diagnosis, and other imbalanced classification problems.
• Used consistently across academic research, industry projects, and ML competitions.

16. Interview-Oriented Points

• Be ready to explain what test_size and random_state control in train_test_split().
• Understand why stratification matters, especially for imbalanced classification problems.
• Be able to explain why train-test splitting should generally happen before scaling/imputation, not after.

17. Exam-Oriented Points

• train_test_split(X, y, test_size=..., random_state=..., stratify=...) is the standard Scikit-learn function for this task.
• test_size controls the proportion reserved for testing; random_state ensures reproducibility.
• stratify=y preserves class proportions between training and test sets, important for imbalanced classification.

18. Comparison Table — traintestsplit vs Cross-Validation (Preview)

19. Quick Revision

• train_test_split() divides features (X) and labels (y) into training and testing subsets.
• test_size controls the proportion for testing; random_state ensures reproducibility.
• stratify=y preserves class proportions, especially important for imbalanced classification problems.
• Perform the split before fitting scalers/imputers to avoid data leakage from the test set.