31. Train Test Split in Machine Learning
31.3 Stratify and random_state Explained Simply
Why. Two small arguments save you from silent bias: stratify and random_state.
What.
- random_state — fixes the shuffle so the same split repeats. Good for teaching and debugging.
- stratify=y — keeps the same yes/no ratio in train and test. Useful when late orders are rare (e.g. only 10% late).
How.
import pandas as pd
from sklearn.model_selection import train_test_split
df = pd.DataFrame({
"x": range(20),
"was_late": [0]*18 + [1, 1], # rare late cases (fictional)
})
X = df[["x"]]
y = df["was_late"]
X_tr, X_te, y_tr, y_te = train_test_split(
X, y, test_size=0.25, random_state=42, stratify=y
)
print("train late rate", y_tr.mean())
print("test late rate", y_te.mean())
What you should see. Train and test late rates stay close. Without stratify, a tiny dataset might put both late rows only in train — then test never sees a late case.
Ravindra Bagale's Tip
Classification + rare class asel tar stratify almost always. Regression (number predict) madhe stratify usually nahi. Fark lakshat theva.