Стратифицированный поезд / тест-сплит в scikit-learn

88

Мне нужно разделить мои данные на обучающий набор (75%) и тестовый набор (25%). В настоящее время я делаю это с помощью кода ниже:

X, Xt, userInfo, userInfo_train = sklearn.cross_validation.train_test_split(X, userInfo)   

Однако я хотел бы стратифицировать свой набор тренировочных данных. Как я могу это сделать? Я изучал этот StratifiedKFoldметод, но не позволял мне указывать разделение 75% / 25% и только стратифицировать набор обучающих данных.

пир
источник

Ответы:

153

[обновление для 0.17]

См. Документы sklearn.model_selection.train_test_split:

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    stratify=y, 
                                                    test_size=0.25)

[/ обновление для 0.17]

Здесь есть запрос на перенос . Но вы можете просто сделать train, test = next(iter(StratifiedKFold(...))) и использовать индексы обучения и тестирования, если хотите.

Андреас Мюллер
источник
1
@AndreasMueller Есть ли простой способ стратифицировать данные регрессии?
Jordan
3
@Jordan в scikit-learn ничего не реализовано. Я не знаю стандартного способа. Мы могли бы использовать процентили.
Андреас Мюллер
@AndreasMueller Вы когда-нибудь видели поведение, при котором этот метод значительно медленнее, чем StratifiedShuffleSplit? Я использовал набор данных MNIST.
snymkpr
@activatedgeek, что кажется очень странным, поскольку train_test_split (... stratify =) просто вызывает StratifiedShuffleSplit и берет первый сплит. Не стесняйтесь открывать вопрос на трекере с воспроизводимым примером.
Андреас Мюллер
@AndreasMueller На самом деле я не открывал проблему, потому что у меня есть сильное чувство, что я делаю что-то не так (хотя это всего 2 строки). Но если я все еще смогу воспроизвести это сегодня несколько раз, я сделаю это!
snymkpr
29

TL; DR: используйте StratifiedShuffleSplit сtest_size=0.25

Scikit-learn предоставляет два модуля для стратифицированного разбиения:

  1. StratifiedKFold : этот модуль полезен в качестве прямого k-кратного оператора перекрестной проверки: в нем настраиваются наборы для n_foldsобучения / тестирования, так что классы одинаково сбалансированы в обоих.

Вот код (прямо из документации выше)

>>> skf = cross_validation.StratifiedKFold(y, n_folds=2) #2-fold cross validation
>>> len(skf)
2
>>> for train_index, test_index in skf:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
...    #fit and predict with X_train/test. Use accuracy metrics to check validation performance
  1. StratifiedShuffleSplit : этот модуль создает единый набор для обучения / тестирования, имеющий одинаково сбалансированные (стратифицированные) классы. По сути, это то, что вы хотите от n_iter=1. Здесь вы можете указать размер теста так же, как вtrain_test_split

Код:

>>> sss = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=0)
>>> len(sss)
1
>>> for train_index, test_index in sss:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
>>> # fit and predict with your classifier using the above X/y train/test
острый
источник
5
Обратите внимание, что по состоянию на 0.18.x, n_iterдолжно быть n_splitsдля StratifiedShuffleSplit - и что для него есть немного другой API: scikit-learn.org/stable/modules/generated/…
lollercoaster
2
Если yэто серия Pandas, используйтеy.iloc[train_index], y.iloc[test_index]
Owlright 01
1
@Owlright Я попытался использовать фрейм данных pandas, и индексы, которые возвращает StratifiedShuffleSplit, не являются индексами в фрейме данных. dataframe index: 2,3,5 the first split in sss:[(array([2, 1]), array([0]))]:(
Meghna Natraj
2
@tangy, почему это цикл for? Разве это не тот случай, когда при X_train, X_test = X[train_index], X[test_index]вызове строки она отменяет X_trainи X_test? Почему тогда не сингл next(sss)?
Bartek Wójcik
13

Вот пример непрерывных / регрессионных данных (пока эта проблема на GitHub не будет решена).

min = np.amin(y)
max = np.amax(y)

# 5 bins may be too few for larger datasets.
bins     = np.linspace(start=min, stop=max, num=5)
y_binned = np.digitize(y, bins, right=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, 
    y, 
    stratify=y_binned
)
  • Где startmin и stopmax вашей непрерывной цели.
  • Если вы не установите, right=Trueтогда ваше максимальное значение будет более или менее выделено в отдельную корзину, и ваше разделение всегда будет терпеть неудачу, потому что в этой дополнительной корзине будет слишком мало образцов.
Иордания
источник
12

Вы можете просто сделать это с помощью train_test_split()метода, доступного в Scikit learn:

from sklearn.model_selection import train_test_split 
train, test = train_test_split(X, test_size=0.25, stratify=X['YOUR_COLUMN_LABEL']) 

Я также подготовил небольшой GitHub Gist, который показывает, как stratifyработает опция:

https://gist.github.com/SHi-ON/63839f3a3647051a180cb03af0f7d0d9

Шаян Амани
источник
6

В дополнение к принятому ответу @Andreas Mueller просто хочу добавить это как @tangy, упомянутый выше:

StratifiedShuffleSplit больше всего похож на train_test_split (stratify = y) с дополнительными функциями:

  1. стратифицировать по умолчанию
  2. указав n_splits , он многократно разбивает данные
Максимум
источник
0
#train_size is 1 - tst_size - vld_size
tst_size=0.15
vld_size=0.15

X_train_test, X_valid, y_train_test, y_valid = train_test_split(df.drop(y, axis=1), df.y, test_size = vld_size, random_state=13903) 

X_train_test_V=pd.DataFrame(X_train_test)
X_valid=pd.DataFrame(X_valid)

X_train, X_test, y_train, y_test = train_test_split(X_train_test, y_train_test, test_size=tst_size, random_state=13903)
Хосе Карлос Кастро
источник
0

Обновление ответа @tangy сверху до текущей версии scikit-learn: 0.23.2 ( документация StratifiedShuffleSplit ).

from sklearn.model_selection import StratifiedShuffleSplit

n_splits = 1  # We only want a single split in this case
sss = StratifiedShuffleSplit(n_splits=n_splits, test_size=0.25, random_state=0)

for train_index, test_index in sss.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
Рой Бахуми
источник