Что такое стратификация?

Стратификация в машинном обучении

Стратификация — это техника разделения данных на тренировочное и тестовое подмножества таким образом, чтобы распределение целевой переменной (класса) было одинаковым в обоих наборах. Она особенно важна при работе с несбалансированными классами.

Проблема без стратификации

from sklearn.model_selection import train_test_split
import numpy as np

# Несбалансированный датасет
y = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1])  # 90% класс 0, 10% класс 1

# Обычное разделение (БЕЗ стратификации)
from sklearn.model_selection import train_test_split

X_dummy = np.arange(10).reshape(-1, 1)
X_train, X_test, y_train, y_test = train_test_split(
    X_dummy, y, test_size=0.3, random_state=42
)

print("БЕЗ стратификации:")
print(f"В тренировке: класс 0 = {(y_train == 0).sum()}, класс 1 = {(y_train == 1).sum()}")
print(f"В тесте: класс 0 = {(y_test == 0).sum()}, класс 1 = {(y_test == 1).sum()}")
print(f"Процент класса 1 в тесте: {(y_test == 1).sum() / len(y_test) * 100:.1f}%")
print(f"Процент класса 1 в исходных данных: {(y == 1).sum() / len(y) * 100:.1f}%")

# Результат: распределение может быть совсем другим!

Стратификация решает проблему

# Разделение СО стратификацией
X_train, X_test, y_train, y_test = train_test_split(
    X_dummy, y, test_size=0.3, random_state=42, stratify=y  # ВАЖНО: stratify=y
)

print("\nСО стратификацией:")
print(f"В тренировке: класс 0 = {(y_train == 0).sum()}, класс 1 = {(y_train == 1).sum()}")
print(f"В тесте: класс 0 = {(y_test == 0).sum()}, класс 1 = {(y_test == 1).sum()}")
print(f"Процент класса 1 в тесте: {(y_test == 1).sum() / len(y_test) * 100:.1f}%")
print(f"Процент класса 1 в исходных данных: {(y == 1).sum() / len(y) * 100:.1f}%")

# Результат: распределение одинаковое в обоих наборах!

Стратифицированная K-fold кросс-валидация

from sklearn.model_selection import StratifiedKFold
import pandas as pd

# Создание несбалансированного датасета
X = np.random.randn(100, 5)
y = np.concatenate([np.zeros(90), np.ones(10)])  # 90% класс 0, 10% класс 1

print("Исходное распределение:")
print(f"Класс 0: {(y == 0).sum()}, Класс 1: {(y == 1).sum()}")

# Обычная K-fold
from sklearn.model_selection import KFold

kf = KFold(n_splits=5, shuffle=True, random_state=42)
for fold, (train_idx, val_idx) in enumerate(kf.split(X, y)):
    y_train = y[train_idx]
    y_val = y[val_idx]
    print(f"\nFold {fold+1} (обычная):")
    print(f"  Тренировка - класс 0: {(y_train == 0).sum()}, класс 1: {(y_train == 1).sum()}")
    print(f"  Валидация - класс 0: {(y_val == 0).sum()}, класс 1: {(y_val == 1).sum()}")
    if len(y_val[y_val == 1]) == 0:
        print(f"  ПРОБЛЕМА: класса 1 нет в валидации!")

# Стратифицированная K-fold
print("\n" + "="*50)
print("Стратифицированная K-fold:")
print("="*50)

skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
    y_train = y[train_idx]
    y_val = y[val_idx]
    print(f"\nFold {fold+1} (стратифицированная):")
    print(f"  Тренировка - класс 0: {(y_train == 0).sum()}, класс 1: {(y_train == 1).sum()}")
    print(f"  Валидация - класс 0: {(y_val == 0).sum()}, класс 1: {(y_val == 1).sum()}")
    print(f"  Процент класса 1 в валидации: {(y_val == 1).sum() / len(y_val) * 100:.1f}%")

Стратифицированная кросс-валидация в Sklearn

from sklearn.model_selection import cross_val_score, StratifiedKFold
from sklearn.ensemble import RandomForestClassifier

X = np.random.randn(200, 10)
y = np.concatenate([np.zeros(180), np.ones(20)])  # 90% vs 10%

model = RandomForestClassifier(random_state=42)

# Кросс-валидация со стратификацией
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
scores = cross_val_score(
    model, X, y, cv=skf, scoring='roc_auc'
)

print(f"ROC-AUC scores со стратификацией: {scores}")
print(f"Mean: {scores.mean():.4f}, Std: {scores.std():.4f}")

Стратификация для многоклассовой классификации

from sklearn.datasets import load_iris

iris = load_iris()
X, y = iris.data, iris.target

print("Исходное распределение:")
unique, counts = np.unique(y, return_counts=True)
for u, c in zip(unique, counts):
    print(f"  Класс {u}: {c} примеров ({c/len(y)*100:.1f}%)")

# Стратифицированное разделение
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=42
)

print("\nРазделение тренировка/тест со стратификацией:")
print("\nТренировка:")
unique, counts = np.unique(y_train, return_counts=True)
for u, c in zip(unique, counts):
    print(f"  Класс {u}: {c} примеров ({c/len(y_train)*100:.1f}%)")

print("\nТест:")
unique, counts = np.unique(y_test, return_counts=True)
for u, c in zip(unique, counts):
    print(f"  Класс {u}: {c} примеров ({c/len(y_test)*100:.1f}%)")

Практический пример: несбалансированная классификация

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import precision_recall_curve, auc

# Создание несбалансированного датасета
X, y = make_classification(
    n_samples=1000, n_features=20, n_classes=2,
    weights=[0.95, 0.05],  # 95% класс 0, 5% класс 1
    random_state=42
)

print(f"Распределение классов: {(y == 0).sum()} vs {(y == 1).sum()}")

# Модель с кросс-валидацией
model = LogisticRegression(random_state=42)
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

from sklearn.model_selection import cross_validate
from sklearn.metrics import recall_score, precision_score, f1_score

scoring = {
    'precision': 'precision',
    'recall': 'recall',
    'f1': 'f1'
}

results = cross_validate(
    model, X, y, cv=skf, scoring=scoring
)

print("\nРезультаты со стратифицированной кросс-валидацией:")
for metric in scoring.keys():
    scores = results[f'test_{metric}']
    print(f"{metric.upper()}: {scores.mean():.4f} (+/- {scores.std():.4f})")

Когда использовать стратификацию

1. Несбалансированные классы: когда один класс встречается намного чаще другого
2. Малый датасет: когда важна каждая выборка, включая редкий класс
3. Кросс-валидация: при оценке качества модели
4. Множественные классы: когда распределение неравномерное

Когда стратификация не нужна

• Хорошо сбалансированные классы: почти одинаковое количество каждого класса
• Регрессия: целевая переменная непрерывна
• Очень большой датасет: случайное разделение даст похожее распределение

Типы стратификации в Sklearn

Стратификация — это один из самых важных приёмов при работе с несбалансированными данными, которые встречаются в 90% реальных задач классификации.