Создание выборок в условиях дисбаланса классов

Задачи классификации зачастую характеризуются несбалансированностью классов, когда наблюдения одного типа сильно превалируют над другими. При этом такая ситуация может иметь естественные причины, например, опрашивая 10000 человек для создания выборки о диагностики рака, разумно ожидать, что соотношение заболевших и здоровых будет не равным.

Соответственно, разбиение данных при формировании тренировочной и тестовой выборок должно проводиться с учетом такой несбалансированности. Иначе вы можете обучить модель на датасете из всех здоровых людей и, тем самым, никакой информации о болезни она не получит. Технически это выразится в том, что модель будет склонна всегда прогнозировать, что человек здоров (и для обучения она будет идеальна).

Рассмотрим методы создания выборок с сохранением (по-возможности) баланса классов в тренировочной и тестовой выборках. Для начала создадим датафрейм:

import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold, KFold, train_test_split
np.random.seed(0)

ind_l = ['one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine','ten']
df = pd.DataFrame([[1, 20, 5, 1], [2, 23, 10, 0], [3, 28, 4, 0], [4, 100, 0, 0], [5, 1, 4, 1],
[6, 6, 6, 1], [7, 34, 2, 0], [8, 45, 4, 0], [9, 33, 23, 0], [10, 35,3, 0]],
columns=['id', 'revenue', 'balance', 'target'], index=ind_l)
df

Если воспользоваться обычным методом кросс-валидации - KFold, то получим следующее:

sp = KFold(n_splits=3) 
for tr_idx, ts_idx in sp.split(df, df['target']): 
     df_train, df_test = df.iloc[tr_idx], df.iloc[ts_idx] 
     print(f'тренировочн - {tr_idx}\n тестов - {ts_idx} ')
     print(f"Соотношение 1/0 в train {df_train[df_train['target']==1].shape[0]/df_train[df_train['target']==0].shape[0]}, в test - {df_test[df_test['target']==1].shape[0]/df_test[df_test['target']==0].shape[0]}")

Данный метод не сохраняет пропорции между классами, чтобы это исправить можно воспользоваться StratifiedKFold:

sp = StratifiedKFold(n_splits=3)
for tr_idx, ts_idx in sp.split(df, df['target']):
     df_train, df_test = df.iloc[tr_idx], df.iloc[ts_idx]
     print(f'тренировочн - {tr_idx}\n тестов - {ts_idx} ')
     print(f"Соотношение 1/0 в train {df_train[df_train['target']==1].shape[0]/ \
     df_train[df_train['target']==0].shape[0]}, \
     в test - {df_test[df_test['target']==1].shape[0]/df_test[df_test['target']==0].shape[0]}")

Как видим, StratifiedKFold по-возможности сохраняет пропорции. При этом в первом случае не получилось идеально сбалансировать в силу размера нашей выборки - 10 наблюдений, 3 фолда (в один попадают 4 элемента тестовых, один из них класса - 1).

Если проводить разбиение методом train_test_split, используйте параметр stratify:

X_tr, X_ts, y_tr, y_ts = train_test_split(df, df['target'], shuffle=False, test_size=0.2)

y_tr[y_tr==1].shape[0]/y_tr[y_tr==0].shape[0], y_ts[y_ts==1].shape[0]/y_ts[y_ts==0].shape[0]

X_tr, X_ts, y_tr, y_ts = train_test_split(df, df['target'], stratify=df['target'], test_size=0.2)

y_tr[y_tr==1].shape[0]/y_tr[y_tr==0].shape[0], y_ts[y_ts==1].shape[0]/y_ts[y_ts==0].shape[0]

Во втором случае лучше удалось сохранить соотношение, однако так же из-за размера данных (10), тестовой выборки (2), лучшим вариантом было оставить в тесте одно наблюдение класса 1 и одно класса 0.