Использование sklearn.model_selection для разделения несбалансированного набора данных

Question

Я использую следующий код для разбиения моего набора данных на наборы train / val / test.

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = 
        train_test_split(X_data, y_data, test_size=0.3, random_state=42)

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

Проблема в том, что мой набор данных действительно несбалансирован.Некоторые классы имеют 500 образцов, а некоторые, например, 70.Является ли этот метод расщепления точным в этой ситуации?Является ли выборка случайной или sklearn использует некоторые методы, чтобы сохранить распределение данных одинаковым во всех наборах?

Answer 1

Вы должны использовать опцию stratify (см. документы ):

X_train, X_test, y_train, y_test = 
        train_test_split(X_data, y_data, test_size=0.3, random_state=42, stratify=y_data)

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