train_test_split без рассинхронизации X и y

В прошлом уроке мы загружали датасет diabetes и обучали модель на первых 400 строках, а оставшиеся строки оставались «за бортом» — без случайного перемешивания и без формальной тестовой выборки. Теперь сделаем это правильно: train_test_split одним вызовом случайно перемешает строки и разобьёт и X, и y синхронно — так, что ни одна тестовая строка не попадёт в обучение.

train_test_split: один вызов для X и y, параметры test_size и random_state

Функция train_test_split принимает произвольное количество массивов или DataFrame-ов и разбивает каждый из них одинаково — с одной и той же случайной перестановкой строк. Именно поэтому X и y передаются вместе:

from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split

bunch = load_diabetes(as_frame=True)
X = bunch.data
y = bunch.target

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=0
)

print(X_train.shape, X_test.shape)
print(y_train.shape, y_test.shape)

Датасет diabetes содержит 442 строки. test_size=0.2 означает, что 20 % строк попадут в тест. Scikit-learn округляет дробное число строк вверх, поэтому в тест уйдут 89 строк, в обучение останутся 353. Точное число всегда можно проверить через X_test.shape[0] — ручные расчёты здесь лишние.

random_state=0 фиксирует генератор случайных чисел. Функция внутри перемешивает индексы строк перед нарезкой, и конкретный порядок этого перемешивания определяется seed-ом. Запустите код дважды с random_state=0 — получите идентичные X_train и X_test. Если не передать random_state (или передать None), функция использует текущее состояние глобального генератора: разбиение обычно меняется при каждом новом запуске, и воспроизвести его позже неудобно. Фиксированное целое число делает пример повторяемым без лишних усилий.

Функция возвращает результаты поочерёдно для каждого переданного массива: сначала X_train и X_test для X, затем y_train и y_test для y. Это и есть правильный порядок распаковки в примере выше. Если передать три массива, получите шесть переменных в том же принципе. Важно распаковать их в правильном порядке — иначе X_test окажется там, где ожидался y_train, и никакой ошибки Python не бросит.

stratify: сохранить доли классов при разделении

Для классификации случайный split можно сделать с учётом меток. Здесь 80 объектов класса 0 и 20 класса 1. В тест откладываем 20 объектов:

from collections import Counter
from sklearn.model_selection import train_test_split

X = [[i] for i in range(100)]
y = [0] * 80 + [1] * 20
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(Counter(y_train))
print(Counter(y_test))

В обучении будет 64 объекта класса 0 и 16 класса 1, в тесте — 16 и 4. stratify=y поддерживает соотношение классов; он не выравнивает их до 50/50.

Если у класса всего один объект, стратифицированное разделение закончится ошибкой. Проверяйте Counter(y) и размеры частей до split: в каждой части должно хватать мест для представления классов. stratify нельзя сочетать с shuffle=False; для временной проверки это другой способ разделения, а не повод перемешать даты ради отсутствия ошибки.