Введение

Когда речь заходит о таких инструментах, как Airflow, MLflow или Docker, многие сразу представляют себе продакшен-среду, и новичков это может пугать. Однако на самом деле эти инструменты полезны не только в проде или крупных компаниях.

Сегодня я хочу рассказать об MLflow. Эта статья рассчитана на тех, кто только начинает свой путь в машинном обучении и обладает базовыми знаниями, а также на практикующих ученых в области ИИ, которые пока не знакомы с этим инструментом или сознательно им не пользуются.

Хотя существуют и альтернативные инструменты (например, DVC или Weights&Biases), MLflow остаётся востребованным решением благодаря простоте освоения и универсальности (подходит как для начинающих, так и для опытных специалистов).

Я не буду давать определения сущностей и объектов, которые есть в MLflow. Лучше всего про это почитать документацию, она сделана достойна. Здесь же мы сосредоточимся на базовом функционале, как этот инструмент может быть полезен.

Github с кодом

tg: @kirill_zakharov_blog

Мотивация

Итак, зачем вообще нужен MLflow? У этого инструмента есть несколько функций, и мы рассмотрим не все из них. Однако ключевые — управление версиями ML-моделей и логирование экспериментов.

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

Особенно становится сложно отслеживать эксперименты, когда проект растет, и тут важно иметь историю экспериментов для сравнения "до" и "после". Также есть человеческий фактор. Даже через неделю бывает сложно вспомнить, какие именно эксперименты уже проводились. А при работе в команде эта проблема многократно усугубляется. Без системы логирования ценные инсайты легко теряются, а время тратится на повторение уже сделанных экспериментов.

В самом простом варианте я бы сказал MLflow — это записная книжка, куда вы вносите записи об экспериментах. Раньше в науке это делали вручную и это называлось “Книга экспериментов”, “Философская тетрадь” или “Дневник наблюдений”. Так давайте освоим эту тетрадь!

Задача

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

satisfaction
neutral or dissatisfied    58879
satisfied                  45025
Name: count, dtype: int64

Всего у нас два класса, поэтому это бинарная классификация и оба класса хорошо сбалансированы.

Мы не будем вдаваться в подробности самой модели, так как цель не в этом, но в процессе построения будем варьировать разные гиперпараметры, с целью выбора наилучшей модели. Шаги с обработкой данных я пропускаю. Итоговый датасет имеет 27 признаков и один столбец — таргет.

Обучать будем две модели из классического ML: логистическую регрессию и градиентный бустинг из библиотеки XGBoost, а также нейронную сеть при помощи Pytorch.

Гиперпараметры для запусков эксперимента удобнее всего хранить в словаре.

params = {
    "max_depth": 3,
    "n_estimators": 100,
    "random_state": 22,
    "learning_rate": 1e-2
}

В качестве метрики возьмем только ROC AUC, хотя MLflow позволяет логировать множество метрик.

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

Функция для визуализации матрицы ошибок

def draw_confusion_matrix(prediction):
    cm = confusion_matrix(y_test, prediction, normalize='all')
    fig, axs = plt.subplots(figsize=(5, 5))
    ConfusionMatrixDisplay(cm, display_labels = [0, 1]).plot(cmap='viridis', ax=axs)
    plt.tight_layout()
    plt.close(fig)
    return fig
  
fig1 = draw_confusion_matrix(prediction)

Запуск MLflow

Для работы с MLflow нам понадобится установить всего одну библиотеку (предполагается что библиотеки для ML у вас уже установлены, а именно Scikit-learn, pandas, Pytorch).

pip install mlflow

Сначала стартанем сервер mlflow через терминал. В продакшене эти вещи запускаются на выделенных серверах, но для персонального использования воспользуемся localhost. Для этого выполним команду в терминале.

mlflow server --host 127.0.0.1 --port 8080.

Здесь надо указать локальный хост и порт на котором запустить UI. Запускать эту команду стоит сразу в той директории, где лежит ваш проект или где он будет. Так как при запуске сервера в текущей директории создаются папки mlruns и mlartifacts, которые будут хранить локально информацию об экспериментах.

В браузере вбейте localhost:8080 и вы должны увидеть что-то наподобие этого.

Сейчас у нас пустой список экспериментов, но совсем скоро мы его заполним.

Что можно логировать?

Давайте поговорим о том, что мы вообще можем логировать. Главные методы в MLflow:

mlflow.log_params(dict). Логирует параметры (тот самый словарь params, который мы задали в начале). В MLflow есть аналоги методов в единственном числе, например, log_param() — это позволяет логировать только один параметр.
mlflow.log_metrics(dict). Логирует метрики. В нашем случае метрика одна — ROC AUC, поэтому вызываем метод в единственном числе. Иначе надо просто передать словарь метрик по аналогии с параметрами.
mlflow.set_tag('Model type', 'value'). Задает произвольный тег. Если вам необходимо передать дополнительную информацию о запуске или об эксперименте, то используйте теги чтобы быстро найти нужную вам запись. В примере я задал тег Model type с разными значениями для логистической регрессии и градиентного бустинга, чтобы затем сгруппировать по этому тегу и сравнить результаты моделей между собой.
mlflow.log_figure(figure, 'file_name.png'). Логирует изображения, которые сохранены в памяти, ссылаясь на переменную в которой лежит нарисованный график. Вторым аргументом данный метод принимает название файла, которое будет отображаться в артефактах.
mlflow.log_artifact(path). Более универсальный метод логирования. Позволяет сохранить все файлы и директории — артефакты, которые располагаются по пути path.
mlflow.infer_signature(X_train, y_test). Позволяет определить размеры тензоров на входе и выходе модели. Также эти размеры можно задать вручную, но в данном случае удобнее передать наши примеры данных.
mlflow.sklearn.log_model(model, path,...). Данный метод логирует модель. Для этого надо указать какая именно у нас библиотека (sklearn, xgboost, pytorch), а в аргументах указать обученную модель и путь к артефактам где сохранить. Также опционально можно указать размерности и входные данные.
Здесь отдельно хочу отметить про сохранение нейронных сетей. Иногда мы хотим сохранять чекпоинты нашей сети, для этого надо использовать log_artifact() во время обучения с указанием соответствующего пути, включающего, например, номер эпохи. А уже обученную модель логировать через log_model() после обучения. Если экспериментов много, то аккуратно выбирайте что хранить и сколько, так как все артефакты будут хранится у вас в папке mlartifacts, в данном случае локально.

Логирование экспериментов

Обучим нашу модель градиентного бустинга с разными гиперпараметрами (например, количество деревьев: 10, 50, 100) и выполним логирование для каждого запуска.

Для работы с MLflow импортируем библиотеку.

import mlflow
from mlflow.models import infer_signature

Инициализируем путь где мы трекаем эксперименты и зададим имя эксперимента.

mlflow.set_tracking_uri(uri="http://127.0.0.1:8080")
mlflow.set_experiment("Classic_ML")

Кстати если название совпадает с экспериментом, который вы удалили, но не очистили корзину, то надо запустить в терминале команду

rm -rf mlruns/.trash/*

непосредственно в текущей директории. А если вы работаете с одним и тем же экспериментом, то у вас просто добавятся в нем новые запуски.

Само логирование осуществляется следующим образом.

Запуск логирования

with mlflow.start_run(description='Experiments with *classical* ML algorithms') as run:
    mlflow.log_params(params)
    mlflow.log_metric("roc_auc", rocauc)
    mlflow.log_figure(fig1, 'Confusion_matrix.png')
    
    
    mlflow.set_tag("Model type", "XGBoost")
    signature = infer_signature(X_train, xgb.predict(X_train))


    model_info = mlflow.xgboost.log_model(
        xgb_model=xgb,
        artifact_path="xgb_airplane_data",
        signature=signature,
        input_example=X_train,
        # registered_model_name="XGB_v0",
    )

Эта команда стартанет один запуск (run) в рамках созданного эксперимента. Каждый раз при ее вызове будет создаваться новый run. Все функции мы разобрали выше, здесь только переданы конкретные значения.

Обновим страничку в браузере и увидим наш эксперимент и несколько запусков по нему (runs).

В меню можно выбрать какие столбцы отобразить (я выбрал нашу метрику и тег). Как видно, имя запуска генерируется автоматически, но его можно передать в качестве аргумента mlflow.start_run(run_name='NAME’) и там же задать описание, которое поддерживает markdown разметку. Также в меню можно выбрать сортировку и группировку, но об этом чуть позже.

Перейдя по первому запуску (у меня это thundering-sloth-881) отображается следующая информация. Здесь наши гиперпараметры и значение метрики на тестовой выборке.

Там же, во вкладке артефакты, хранится наша модель, файл requirements с библиотеками и также наша визуализация матрицы ошибок и все что мы логируем через log_artifacts().

Теперь обучим логистическую регрессию и повторим логирование. Я варьировал max_iter: 50, 70, 100.

Логирование логистической регрессии

with mlflow.start_run(description='Experiments with *classical* ML algorithms') as run:
    mlflow.log_params(params)
    mlflow.log_metric("roc_auc", rocauc)
    mlflow.log_figure(fig1, 'Confusion_matrix.png')
    
    mlflow.set_tag("Model type", "LogReg")
    signature = infer_signature(X_train, lr.predict(X_train))
    
    model_info = mlflow.sklearn.log_model(
        lr,
        artifact_path="log_reg_airplane_data",
        signature=signature,
        input_example=X_train,
        # registered_model_name="LogR_v0",
    )

Именно здесь воспользуемся группировкой по нашему тегу Model type и отсортируем запуски в каждой группе по значению метрики. Выглядит это примерно так.

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

Если мы варьируем наши запуски по seed или по разным разбиениям датасета для проверки статистически значимого результата нашей модели, то правильнее будет выбирать среднее значение. Все зависит от целей эксперимента.

И последнее, давайте выберем все запуски и нажмем на кнопку Compare.

Откроется следующая визуализация для всех запусков в нашем эксперименте.

Здесь прекрасно видны наши изменения по гиперпараметрам и значение метрики. Можно поиграться с разными графиками и фильтрами.

Логирование нейронных сетей

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

Инициализация модели и функция обучения

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, TensorDataset
from torchsummary import summary

from tqdm import tqdm

class Classifier(nn.Module):
    def __init__(self, **kwargs):
        super(Classifier, self).__init__()
        
        self.model = nn.Sequential(
            nn.Linear(kwargs.get('in_dim'), 2**7),
            nn.SiLU(),
            nn.Linear(2**7, 2**7),
            nn.SiLU(),
            nn.Linear(2**7, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        out = self.model(x)
        
        return out


NUM_EPOCHS = 100
BATCH_SIZE = 64
torch.manual_seed(55555)

X_train_tensor, X_test_tensor, y_train_tensor, y_test_tensor = \
    torch.tensor(X_train.astype('float32')),\
    torch.tensor(X_test.astype('float32')),\
    torch.tensor(y_train.astype('float32').reshape(-1, 1)),\
    torch.tensor(y_test.astype('float32').reshape(-1, 1))
    
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)


def train(num_epochs):
    model = Classifier(in_dim=27)
    optimizer = optim.AdamW(model.parameters(), lr=1e-3)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=2e-4)

    epochs = tqdm(range(num_epochs))
    criterion = nn.BCELoss()
    
    model.train()

    for i in epochs:
        epoch_loss = 0
        for data, target in train_loader:
            
            output = model(data)
            loss = criterion(output, target)
            
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1)
            optimizer.step()
            
            epoch_loss += loss.item()
        
        scheduler.step()
        
        epoch_loss = epoch_loss/data.size(0)
            
        if i % 10 == 0:
            epochs.set_description(f'Loss: {epoch_loss}')
            mlflow.log_metric("Loss", f"{epoch_loss:3f}")
            
    return model

Создадим новый эксперимент.

mlflow.set_experiment("NeuralNetworks")

Для логирования сделаем все тоже самое, добавив только артефакт с архитектурой сети. Я использую summary из библиотеки torchsummary. Она похожа на реализацию в keras, там сразу указываются размерности тензоров на выходе из каждого слоя, число параметров и память.

Также хочу отметить, что в MLflow реализовано автологирование для Pytorch. Это упрощает некоторые действия, но если хотите гибкости, то используйте разобранные функции.

Обучение нейронной сети с логирование выглядит следующим образом.

Логирование нейронной сети

with mlflow.start_run(description='Experiments with DL algorithms', tags={'version': 'v0'}, run_name='NN_64'):
    mlflow.set_tag("Model type", "Neural Network")
    
    params = {
        "epochs": NUM_EPOCHS,
        "learning_rate_start": 1e-3,
        "batch_size": BATCH_SIZE,
        "loss_function": nn.BCELoss(),
        "optimizer": "AdamW",
        "scheduler": 'optim.lr_scheduler.CosineAnnealingLR()'
    }
    
    mlflow.log_params(params)

    model_NN = train(NUM_EPOCHS)
    
    
    with open("model_summary.txt", "w") as file:
        file.write(str(summary(model_NN, input_size=(BATCH_SIZE, X_train_tensor.size(1)))))
        
    mlflow.log_artifact("model_summary.txt")
    
    
    signature = infer_signature(X_train, model_NN(X_train_tensor).detach().cpu().numpy())
    mlflow.pytorch.log_model(model_NN, "model_NN", signature=signature)
    
    
    model_NN.eval()
    with torch.no_grad():
        prediction_NN = model_NN(X_test_tensor).detach().cpu().numpy().T[0]
        prediction_NN = np.array([1 if i > 0.5 else 0 for i in prediction_NN])
    
    rocauc = roc_auc_score(y_test, prediction_NN)
    mlflow.log_metric("roc_auc", rocauc)
    
    fig1 = draw_confusion_matrix(prediction_NN)
    mlflow.log_figure(fig1, 'Confusion_matrix.png')

Я варьировал размер батча: 32, 64, 128. Теперь выберем в левом верхнем углу сразу все эксперименты и перейдем в режим графика.

Лучшая модель — нейронная сеть, а ее значение метрики 0.94. Конечно в этом туториале мы делали просто хаотичные эксперименты, но тем не менее даже так, теперь мы храним всю информацию о запусках и гиперпараметрах и в нужный момент можем вытащить веса модели.

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

Здесь видно, что мы меняли размер батча, при этом постоянными оставались оптимизатор, количество эпох и т. д.

А под графиком подробно указывается вся информация о запусках.

На этом наш туториал завершается. Мы рассмотрели основы трекинга и логирования ML экспериментов с использованием MLflow. Для более подробного изучения вновь отсылаю к документации, а полный код на моем гите. Делитесь в комментариях своим опытом организации ML-экспериментов.

MLflow для исследований: как систематизировать ML-эксперименты