Попытка разделить мое datafame на репрезентативный набор поездов и тестов

Question

У меня есть некоторые данные. Я хочу разделить его на поезд и испытательный комплекс. Поезд должен содержать 2/3 данных. Я хочу, чтобы два набора были репрезентативными для всего набора. В моем столбце «Класс» у меня есть 4 или 2 для представления двух классов. Я хочу, чтобы мой тестовый набор имел такое же соотношение 4: 2. Для этого я создал следующий фрагмент кода:

trainTotal = 455
benTotal = 296
malTotal = 455-296
b = 0
m = 0
tr = 0
i = 0
j = 0

for index, row in data.iterrows():
    if row['Class'] == 2:
        if tr < trainTotal and b < benTotal:
            train.loc[i] = data.iloc[index]
            b = b+1
            tr = tr + 1
            i = i+1
        else:
            test.loc[j] = data.iloc[index]
            j = j+1
    if row['Class'] == 4:
        if tr < trainTotal and m < malTotal:
            train.loc[i] = data.iloc[index]
            tr = tr + 1            
            i = i + 1
            m = m+1
        else:
            test.loc[j] = data.iloc[index]
            j = j + 1

Я получаю правильное итоговое число значений в моем фрейме данных поезда, но случаи не представлены, как я надеялся. Это входит в if tr < trainTotal and b < benTotal: слишком много раз. Есть идеи, в чем проблема?

Answer 1

Как сказал Майкл Гарднер, train_test_split - это функция, которую вы ищете.

По умолчанию она будет делиться случайным образом, но вы можете использовать stratify, чтобы сообщить ей, что вы хотите такое же соотношениедля вашего столбца класса в наборах данных train и test.

Работает так:

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    data,
    target,
    test_size = 0.3,
    stratify=data[['your_column']]
)