Укажите максимальное количество предметов в классе при использовании train_test_split от sklearn

Question

Я работаю с сильно несбалансированным набором данных и использую train_test_split из sklearn.model_selection

У меня есть 10000 элементов в этом наборе данных, и соотношение составляет примерно 10/2/2/1 для классов, и я ищу способ сбалансировать разделение train и я хотел бы прекратить добавлять элементы в самый большой класс, когда он достигает максимального числа.

Можно ли ограничить количество элементов, я знаю, что можно удалить лишние элементы после разделения, но мне интересно, есть ли такая опция?

Answer 1

Используйте параметр stratify при вызове функции train_test_split. Следуйте документации для получения дополнительной информации.

Для 30% тестовых данных вы можете сделать вот так

X_train,X_test, y_train, y_test = train_test_split(data, y_true, stratify=y_true, test_size=0.3)

data - это ваши общие данные, а y_true - ваши основные значения истинности