Как сделать препроцессинг надежнее и сократить код в три раза

Одним из ключевых этапов подготовки датасета к подаче в модель машинного обучения является кодирование и масштабирование признаков. Рассмотрим, как это сделать надежнее и проще всего. Для демонстрации используем данные о пассажирах Титаника:

import pandas as pd
import numpy as np
np.random.seed(0)

df = pd.read_csv("https://s3.amazonaws.com/h2o-public-test-data/smalldata/gbm_test/titanic.csv")
df['sex'] = df['sex'].map({'male':1, 'female':0})

df[['age', 'body', 'fare']] = df[['age', 'body', 'fare']].fillna(df[['age', 'body', 'fare']].agg('median').to_dict())
df = df.fillna('unknown')

df.head()

Разделим колонки на выборки и поделим по типам:

id_cols = ['name', 'ticket']
target_col = 'survived'

cat_cols = df.drop([target_col]+id_cols, axis=1).select_dtypes(exclude=np.number).columns.tolist()
num_cols = df.drop([target_col]+id_cols, axis=1).select_dtypes(include=np.number).columns.tolist()

display(cat_cols)
display(num_cols)

from sklearn.model_selection import train_test_split

df_tr_val, df_ts = train_test_split(df, test_size=0.2)
df_tr, df_val = train_test_split(df_tr_val, test_size=0.25)
print(df_tr.shape, df_val.shape, df_ts.shape)

Теперь рассмотрим применение энкодера. Обучать его будем только на тренировочной выборке, а применять трансформацию ко всем, как и в реальной ситуации, когда вы будете иметь только обучающий датасет, и вам будет требоваться все этапы преобразований перенести на набор для прогноза. В нашей ситуации этот новый набор имитируют валидационная и тестовая выборки:

from sklearn.preprocessing import OneHotEncoder

ohe = OneHotEncoder(sparse_output=False, drop='first', handle_unknown='ignore')
ohe.fit(df_tr[cat_cols])

В данном случае рассмотрим применение OneHotEncoder-а как одного из ключевых и сложных ввиду появления новых колонок. Чтобы склеить данные в исходные наборы надо проследить за соответствием индексов. При этом имена новых колонок получаются из метода get_feature_names_out (get_feature_names в старых версиях sklearn):

data_tr0 = pd.DataFrame(ohe.transform(df_tr[cat_cols]),
                    columns=ohe.get_feature_names_out(), index=df_tr.index)
data_tr0 = df_tr.drop(columns=cat_cols).join(data_tr0)

data_val0 = pd.DataFrame(ohe.transform(df_val[cat_cols]),
                    columns=ohe.get_feature_names_out(), index=df_val.index)
data_val0 = df_val.drop(columns=cat_cols).join(data_val0)

data_ts0 = pd.DataFrame(ohe.transform(df_ts[cat_cols]),
                    columns=ohe.get_feature_names_out(), index=df_ts.index)
data_ts0 = df_ts.drop(columns=cat_cols).join(data_ts0)

В примере выше мы для ohe части набора создавали датафрейм и задавали тот же индекс, что и в исходном (например, index=df_tr.index). Альтернативно можно сбросить индекс у той части, к которой мы присоединяем ohe признаки:

data_tr1 = pd.DataFrame(ohe.transform(df_tr[cat_cols]), columns=ohe.get_feature_names_out())
data_tr1 = df_tr.drop(cat_cols, axis=1).reset_index(drop=True).join(data_tr1)

data_val1 = pd.DataFrame(ohe.transform(df_val[cat_cols]), columns=ohe.get_feature_names_out())
data_val1 = df_val.drop(cat_cols, axis=1).reset_index(drop=True).join(data_val1)

data_ts1 = pd.DataFrame(ohe.transform(df_ts[cat_cols]), columns=ohe.get_feature_names_out())
data_ts1 = df_ts.drop(cat_cols, axis=1).reset_index(drop=True).join(data_ts1)

А вот еще третий вариантов того же кодирования (без мук с индексами):

data_tr2 = df_tr.drop(columns=cat_cols).copy()
data_tr2[ohe.get_feature_names_out()] = ohe.transform(df_tr[cat_cols])

data_val2 = df_val.drop(columns=cat_cols).copy()
data_val2[ohe.get_feature_names_out()] = ohe.transform(df_val[cat_cols])

data_ts2 = df_ts.drop(columns=cat_cols).copy()
data_ts2[ohe.get_feature_names_out()] = ohe.transform(df_ts[cat_cols])

В идентичности методов легко убедиться способом, о котором я писал ранее:

display(data_tr0.merge(data_tr1, how='outer', on=data_tr0.columns.tolist(), indicator='ind').loc[lambda x:x['ind']!='both'])

С кодированием проще, так как для этого можно просто перезаписать значения колонок:

from sklearn.preprocessing import StandardScaler

sc = StandardScaler()
sc.fit(data_tr0[num_cols])
data_tr0[num_cols] = sc.transform(data_tr0[num_cols])
data_val0[num_cols] = sc.transform(data_val0[num_cols])
data_ts0[num_cols] = sc.transform(data_ts0[num_cols])

Применение трансформатора

С трансформатором колонок то же делается в разы проще. Достаточно задать преобразователи и колонки, на которые распространяется их действие (подробнее читай здесь). Обучение и применение происходят по одинаковой схеме:

from sklearn.compose import make_column_transformer
from sklearn.pipeline import Pipeline

ct = make_column_transformer((StandardScaler(), num_cols),
                             (OneHotEncoder(drop='first', sparse_output=False, handle_unknown='ignore'), cat_cols),
                             remainder='passthrough', verbose_feature_names_out=False)

ct.fit(df_tr)

А теперь применение:

data_tr = pd.DataFrame(ct.transform(df_tr), columns= ct.get_feature_names_out())
data_val = pd.DataFrame(ct.transform(df_val), columns= ct.get_feature_names_out())
data_ts = pd.DataFrame(ct.transform(df_ts), columns= ct.get_feature_names_out())

display(data_ts.merge(data_ts0, how='outer', on=data_tr.columns.tolist(), indicator='ind').loc[lambda x:x['ind']!='both'])