Введение

В этой статье мы рассмотрим, как создать разделение набора данных для обучения и тестирования, сохраняя при этом сбалансированное распределение категорий. Мы будем использовать набор данных CooperUnion, который представляет собой набор данных об автомобилях, включая их марку, модель, год выпуска и различные характеристики. Разделив набор данных на наборы для обучения и тестирования, мы можем оценить производительность модели на невидимых данных. Однако важно поддерживать распределение категорий как в обучающем, так и в тестовом наборах, чтобы наши модели могли хорошо обобщать. Для этого мы будем использовать ключевое слово stratify в функции train_test_split из библиотеки scikit-learn.

Давайте начнем!

Шаг 1. Загрузите и исследуйте набор данных CooperUnion

Во-первых, нам нужно импортировать необходимые библиотеки и загрузить набор данных.

import pandas as pd
from sklearn.model_selection import train_test_split

url = 'https://raw.githubusercontent.com/fenago/datasets/main/CooperUnionDataset.csv'
data = pd.read_csv(url)

data.head()

Набор данных CooperUnion (https://www.kaggle.com/datasets/CooperUnion/cardataset) содержит информацию об автомобилях, такую ​​как их марка, модель, год выпуска, тип топлива в двигателе, мощность двигателя и многое другое. Эти данные можно использовать для создания прогностических моделей для различных целей, таких как прогнозирование цен, эффективности использования топлива или прогнозирование популярности конкретной марки или модели.

Шаг 2. Проверьте распределение категорий

Прежде чем мы разделим данные, давайте рассмотрим распределение категорий. В этом примере мы рассмотрим распределение «make» в наборе данных.

category_distribution = data['Make'].value_counts(normalize=True) * 100
print("Category Distribution:")
print(category_distribution)

Этот код выведет процентное распределение каждой марки в наборе данных.

Шаг 3: обучение-тестирование-разделение без стратификации

Теперь давайте выполним разбиение обучающего теста без стратификации и проверим распределение категорий в полученных наборах данных.

train_data, test_data = train_test_split(data, test_size=0.2, random_state=42)

train_distribution = train_data['Make'].value_counts(normalize=True) * 100
test_distribution = test_data['Make'].value_counts(normalize=True) * 100

print("Train Distribution:")
print(train_distribution)
print("\nTest Distribution:")
print(test_distribution)

На выходе будет показано распределение категорий как в обучающем, так и в тестовом наборах данных, которое может не совпадать с исходным распределением.

Шаг 4: Обучение-тестирование-разделение со стратификацией

Чтобы сохранить одинаковое распределение категорий как в обучающем, так и в тестовом наборах, мы будем использовать ключевое слово stratify в функции train_test_split.

train_data_strat, test_data_strat = train_test_split(data, test_size=0.2, random_state=42, stratify=data['Make'])

train_distribution_strat = train_data_strat['Make'].value_counts(normalize=True) * 100
test_distribution_strat = test_data_strat['Make'].value_counts(normalize=True) * 100

print("Train Distribution (Stratified):")
print(train_distribution_strat)
print("\nTest Distribution (Stratified):")
print(test_distribution_strat)

На выходе будет показано распределение категорий в стратифицированных наборах данных поезда и теста, которое должно быть похоже на исходное распределение.

Заключение

В этой статье мы продемонстрировали, как использовать ключевое слово stratify в функции train_test_split для поддержания распределения категорий как в обучающем, так и в тестовом наборах данных. Этот метод гарантирует, что наша модель обучена и протестирована на данных с аналогичным распределением категорий, что может привести к лучшему обобщению и более точным прогнозам. Это особенно важно при работе с несбалансированными наборами данных, где некоторые категории представлены недостаточно. Используя стратификацию, мы можем гарантировать, что наши модели хорошо приспособлены для работы со всеми категориями, что приводит к более надежным и надежным результатам.

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