Примеры разделения датасета на train и test c Scikit-learn

Если вы разбиваете датасет на данные для обучения и тестирования, нужно помнить о некоторых моментах. Далее следует обсуждение трех передовых практик, которые стоит учитывать при подобном разделении. А также демонстрация того, как реализовать эти соображения в Python.

В данной статье обсуждаются три конкретных особенности, которые следует учитывать при разделении набора данных, подходы к решению связанных проблем и практическая реализация на Python.

Для наших примеров мы будем использовать модуль train_test_split библиотеки Scikit-learn, который очень полезен для разделения датасетов, независимо от того, будете ли вы применять Scikit-learn для выполнения других задач машинного обучения. Конечно, можно выполнить такие разбиения каким-либо другим способом (возможно, используя только Numpy). Библиотека Scikit-learn включает полезные функции, позволяющее сделать это немного проще.

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    train_size=0.67, 
                                                    random_state=42)

Возможно, вы использовали этот модуль для разделения данных в прошлом, но при этом не приняли во внимание некоторые детали.

Случайное перемешивание строк

Первое, на что следует обратить внимание: перемешаны ли ваши экземпляры? Это следует делать пока нет причин не перетасовывать данные (например, они представляют собой временные интервалы). Мы должны убедиться в том, что наши экземпляры не разбиты на выборки по классам. Это потенциально вносит в нашу модель некоторую нежелательную предвзятость.

Например, посмотрите, как одна из версий набора данных iris, упорядочивает свои экземпляры при загрузке:

from sklearn.datasets import load_iris

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

print(f"Классы датасета: {iris.target}")
Классы датасета: [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 2
 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
 2 2]

Если такой набор данных с тремя классами при равном числе экземпляров в каждом разделить на две выборки: 2/3 для обучения и 1/3 для тестирования, то полученные поднаборы будут иметь нулевое пересечение классовых меток. Это, очевидно, недопустимо при изучении признаков для предсказания классов. К счастью, функция train_test_split по умолчанию автоматически перемешивает данные (вы можете переопределить это, установив для параметра shuffle значение False).

  • В функцию должны быть переданы как вектор признаков, так и целевой вектор (X и y).
  • Для воспроизводимости вы должны установить аргумент random_state.
  • Также необходимо определить либо train_size, либо test_size, но оба они не нужны. Если вы явно устанавливаете оба параметра, они должны составлять в сумме 1.

Вы можете убедится, что теперь наши классы перемешаны.

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    train_size=0.67, 
                                                    random_state=42)

print(f"Классы в y_train:\n{y_train}")
print(f"Классы в y_test:\n{y_test}")
Классы в y_train:
[1 2 1 0 2 1 0 0 0 1 2 0 0 0 1 0 1 2 0 1 2 0 2 2 1 1 2 1 0 1 2 0 0 1 1 0 2
 0 0 1 1 2 1 2 2 1 0 0 2 2 0 0 0 1 2 0 2 2 0 1 1 2 1 2 0 2 1 2 1 1 1 0 1 1
 0 1 2 2 0 1 2 2 0 2 0 1 2 2 1 2 1 1 2 2 0 1 2 0 1 2]
Классы в y_test:
[1 0 2 1 1 0 1 2 1 1 2 0 0 0 0 1 2 1 1 2 0 2 0 2 2 2 2 2 0 0 0 0 1 0 0 2 1
 0 0 0 2 1 1 0 0 1 2 2 1 2]

Стратификация (равномерное распределение) классов

Данное размышление заключается в следующем. Равномерно ли распределено количество классов в наборах данных, разделенных для обучения и тестирования?

import numpy as np

print(f"Количество строк в y_train по классам: {np.bincount(y_train)}")
print(f"Количество строк в y_test по классам: {np.bincount(y_test)}")
Количество строк в y_train по классам: [31 35 34]
Количество строк в y_test по классам: [19 15 16]

Это не равная разбивка. Главная идея заключается в том, получает ли наш алгоритм равные возможности для изучения признаков каждого из представленных классов и последующего тестирования результатов обучения, на равном числе экземпляров каждого класса. Хотя это особенно важно для небольших наборов данных, желательно постоянно уделять внимание данному вопросу.

Мы можем задать пропорцию классов при разделении на обучающий и проверяющий датасеты с помощью параметра stratify функции train_test_split. Стоит отметить, что мы будем стратифицировать в соответствии распределению по классам в y.

X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    train_size=0.67, 
                                                    random_state=42,
                                                    stratify=y)

print(f"Количество строк в y_train по классам: {np.bincount(y_train)}")
print(f"Количество строк в y_test по классам: {np.bincount(y_test)}")
Количество строк в y_train по классам: [34 33 33]
Количество строк в y_test по классам: [16 17 17]

Сейчас это выглядит лучше, и представленные числа говорят нам, что это наиболее оптимально возможное разделение.

Дополнительное разделение

Третье соображение относится к проверочным данным (выборке валидации). Есть ли смысл для нашей задачи иметь только один тестовый датасет. Или мы должны подготовить два таких набора — один для проверки наших моделей во время их точной настройки, а еще один — в качестве окончательного датасета для сравнения моделей и выбора лучшей.

train test validation


Если мы определим 2 таких набора, это будет означать, что одна выборка, будет храниться до тех пор, пока все предположения не будут проверены, все гиперпараметры не настроены, а все модели обучены для достижения максимальной производительности. Затем она будет показана моделям только один раз в качестве последнего шага в наших экспериментах.

Если вы хотите использовать датасеты для тестирования и валидации, создать их с помощью train_test_split легко. Для этого мы разделяем весь набор данных один раз для выделения обучающей выборки. Затем еще раз, чтобы разбить оставшиеся данные на датасеты для тестирования и валидации.

Ниже, используя набор данных digits, мы разделяем 70% для обучения и временно назначаем остаток для тестирования. Не забывайте применять методы, описанные выше.

from sklearn.datasets import load_digits

digits = load_digits()
X, y = digits.data, digits.target

X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    train_size=0.7, 
                                                    random_state=42,
                                                    stratify=y)

print(f"Количество строк в y_train по классам: {np.bincount(y_train)}")
print(f"Количество строк в y_test по классам: {np.bincount(y_test)}")
Количество строк в y_train по классам: [124 127 124 128 127 127 127 125 122 126]
Количество строк в y_test по классам: [54 55 53 55 54 55 54 54 52 54]

Обратите внимание на стратифицированные классы в полученных наборах. Затем мы повторно делим тестовый датасет.

X_test, X_val, y_test, y_val = train_test_split(X_test, y_test, 
                                                    train_size=0.5, 
                                                    random_state=42,
                                                    stratify=y_test)

print(f"Количество строк в y_test по классам: {np.bincount(y_test)}")
print(f"Количество строк в y_val по классам: {np.bincount(y_val)}")
Количество строк в y_test по классам: [27 27 27 27 27 28 27 27 26 27]
Количество строк в y_val по классам: [27 28 26 28 27 27 27 27 26 27]

Обратите внимание на стратификацию классов по всем наборам данных, которая является оптимальной.
Теперь вы готовы обучать, проверять и тестировать столько моделей машинного обучения, сколько вы сочтете нужным для ваших данных.

Еще один совет: вы можете подумать об использовании перекрестной валидации вместо простой стратегии обучение/тестирование или обучение/валидация/тестирование. Мы рассмотрим вопросы кросс-валидации в следующий раз.

Подписывайтесь на канал в Дзене

Полезный контент для начинающих и опытных программистов в канале Лента Python разработчика — Как успевать больше, делать лучше и не потерять мотивацию.

Обучение Python и Data Science

Профессия Data Scientist

Профессия Data Scientist

11 520 5 760 ₽/мес.
Профессия Python-разработчик

Профессия Python-разработчик

7 820 3 910 ₽/мес.
Профессия Python Fullstack

Профессия Python Fullstack

7 820 3 910 ₽/мес.
Курс Аналитик данных с нуля

Курс Аналитик данных с нуля

6 500 3 900 ₽/мес.

Появились вопросы? Задайте на Яндекс.Кью

У сайта есть сообщество на Кью >> Python Q <<. Там я, эксперты и участники отвечаем на вопросы по python и программированию.