Train, validation и test: где выбирать модель

Сравнение трёх вариантов по MAE

Для выбора лучшего решения мы заранее определяем трёх кандидатов. Это не какие-то экзотические алгоритмы, а базовые и понятные варианты. Первый — baseline_mean на основе DummyRegressor со стратегией mean. Ему вообще не нужны признаки площади или комнат: при вызове predict он просто возвращает среднее значение арендной платы, посчитанное на обучающей выборке y_train. Это наша точка отсчёта. Второй кандидат — линейная регрессия LinearRegression только на одном признаке area_m2. Третий — та же линейная регрессия, но уже на двух признаках: area_m2 и rooms.

Все три модели проходят через строгий одинаковый регламент: обучаются исключительно на 48 квартирах train, строят прогноз для одних и тех же 16 квартир validation, после чего качество каждого прогноза оценивается функцией mean_absolute_error. Метрика MAE измеряется в тех же единицах, что и целевая переменная — в рублях за месяц аренды. Чем ниже значение ошибки, тем ближе прогнозы к реальным ставкам.

import pandas as pd
from sklearn.dummy import DummyRegressor
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split

data = pd.read_csv("rent.csv", index_col="listing_id")
X = data[["area_m2", "rooms"]]
y = data["rent_rub"]

def split_rows(X, y):
    X_dev, X_test, y_dev, y_test = train_test_split(
        X, y, test_size=0.20, random_state=42)
    X_train, X_val, y_train, y_val = train_test_split(
        X_dev, y_dev, test_size=0.25, random_state=42)
    return X_train, X_val, X_test, y_train, y_val, y_test

X_train, X_val, X_test, y_train, y_val, y_test = split_rows(X, y)
train_ids, val_ids, test_ids = map(set, (X_train.index, X_val.index, X_test.index))
assert train_ids.isdisjoint(val_ids) and train_ids.isdisjoint(test_ids)
assert val_ids.isdisjoint(test_ids)
assert train_ids | val_ids | test_ids == set(X.index)
again = split_rows(X, y)
assert all(a.index.equals(b.index) for a, b in zip(
    (X_train, X_val, X_test, y_train, y_val, y_test), again))
print(f"train={len(X_train)}, validation={len(X_val)}, test={len(X_test)}")

candidates = [
    ("baseline_mean", DummyRegressor(strategy="mean"), ["area_m2"]),
    ("linear_area", LinearRegression(), ["area_m2"]),
    ("linear_area_rooms", LinearRegression(), ["area_m2", "rooms"]),
]
results = []
for name, model, columns in candidates:
    model.fit(X_train[columns], y_train)
    prediction = model.predict(X_val[columns])
    mae = mean_absolute_error(y_val, prediction)
    results.append((name, float(mae)))
    print(f"{name}: validation MAE = {mae:.2f} RUB")
winner = min(results, key=lambda row: row[1])[0]
print(f"Selected: {winner}; test not evaluated")

При запуске этого кода на нашем выданном наборе rent.csv мы получаем следующий вывод консоли:

train=48, validation=16, test=16
baseline_mean: validation MAE = 8550.79 RUB
linear_area: validation MAE = 2241.82 RUB
linear_area_rooms: validation MAE = 1606.18 RUB
Selected: linear_area_rooms; test not evaluated

Простой прогноз по среднему ошибается в среднем на 8550.79 рублей в месяц. Модель, использующая площадь, заметно точнее — её ошибка составляет 2241.82 рубля. Добавление второго признака (количества комнат) снижает среднюю ошибку до 1606.18 рублей. Победитель определяется по наименьшему сырому значению типа float, а округление до двух знаков применяется только при печати строки на экране.