Что такое MLFlow и как он помогает в разработке моделей

Александр Рыжков

Ментор Skillfactory, руководитель команды LightAutoML и 4х Kaggle Grandmaster

MLflow — это инструмент, который стал для меня (и многих команд, с которыми я работал) настоящим спасением от хаоса в ML-проектах. 

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

Что такое MLflow и зачем он нужен?

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

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

Ключевые задачи, которые решает MLflow:

1. Отслеживание экспериментов: запись параметров, метрик, версий кода и артефактов (например, обученных моделей) для каждого запуска. Больше никаких потерянных результатов!

2. Воспроизводимость: упаковка кода и его зависимостей так, чтобы любой мог запустить ваш эксперимент и получить те же результаты (при условии тех же данных, конечно).

3. Управление моделями: централизованное хранилище для ваших обученных моделей с версионированием и управлением стадиями (например, «разработка», «тестирование», «продакшн»).

4. Развертывание моделей: MLflow предлагает инструменты для упрощения развертывания моделей в различных средах.

Думайте об MLflow как об организованном лабораторном журнале, системе контроля версий для моделей и конвейере для их доставки — все в одном флаконе.

MLflow для трекинга экспериментов — основные компоненты

Это, пожалуй, самая часто используемая часть MLflow, особенно на начальных этапах. Компонент mlflow.tracking позволяет логировать все самое важное во время обучения модели. Давайте разберемся с основными понятиями:

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

• Run (Запуск). Один конкретный запуск вашего кода. Например, обучение модели с определенным набором гиперпараметров. Каждый раз, когда вы запускаете скрипт обучения с логированием в MLflow, создается новый Run.

• Parameters (Параметры). Входные данные вашего запуска, которые вы хотите сохранить. Обычно это гиперпараметры модели (скорость обучения, глубина дерева), версии используемых данных, фичи. 

• Metrics (Метрики). Числовые значения, которые показывают качество вашего запуска. Это могут быть точность (accuracy), ошибка (loss), F1-score. Метрики могут меняться со временем (например, loss по эпохам). 

• Artifacts (Артефакты). Любые файлы, которые вы хотите сохранить как результат запуска. Самый частый артефакт — это сама обученная модель. Также это могут быть графики (например, learning curve), файлы с предсказаниями, примеры данных. 

Все это добро удобно просматривать во встроенном веб-интерфейсе MLflow UI. Вы можете запустить его командой mlflow ui в терминале, и он покажет вам все ваши эксперименты, запуски, позволит сравнивать их, сортировать по метрикам и параметрам. Это невероятно удобно для анализа результатов!

Что такое MLflow Model Registry?

Итак, вы провели кучу экспериментов и нашли несколько отличных моделей (артефактов). Что дальше? Как управлять ими? Как понять, какая модель сейчас используется в продакшене, а какая только тестируется? Здесь на помощь приходит MLflow Model Registry (реестр моделей).

Model Registry позволяет:

1. Централизованно хранить ваши лучшие модели.

2. Версионировать модели. Если вы обучили новую, улучшенную версию той же модели, она получит новый номер версии (например, ChurnPredictor v2).

3. Присваивать моделям стадии (Stages):

• Staging: Модель на тестировании, кандидат на переход в продакшен.

• Production: Модель, которая используется в реальной системе.

• Archived: Устаревшая модель, больше не используется.

• Можно добавлять и свои стадии.

4. Добавлять описания и теги к моделям и их версиям.

5. Отслеживать происхождение модели (из какого Run она была создана).

Это делает процесс перехода модели от разработки к продакшену гораздо более контролируемым и прозрачным. Команда DevOps или другие разработчики могут легко понять, какую версию модели им нужно взять для развертывания.

Проекты в MLflow — как использовать? Создание конфига проекта

Компонент MLflow Projects — это стандартный формат для упаковки вашего кода машинного обучения так, чтобы его можно было легко запустить где угодно — на вашей машине, на машине коллеги или в облаке.

Основная идея — описать ваш проект в специальном файле MLproject (в формате YAML). Этот файл содержит:

• name: Имя вашего проекта (необязательно).

• environment: Указание, как создать окружение для запуска кода. Чаще всего используется ссылка на файл conda.yaml, который описывает все необходимые Python-библиотеки и их версии. Также можно использовать Docker-контейнер.

• entry_points: Точки входа — команды, которые можно выполнить в рамках проекта. Обычно есть основная точка входа (main), которая запускает ваш скрипт обучения или предсказания. Здесь же можно определить параметры, которые можно передать в скрипт при запуске.

Пример простого MLproject-файла:

name: tutorial

environment: conda.yaml

entry_points:
  main:
    parameters:
      alpha: {type: float, default: 0.5}
      l1_ratio: {type: float, default: 0.1}
    command: "python train.py --alpha {alpha} --l1_ratio {l1_ratio}"

А вот пример conda.yaml:

name: tutorial-env
channels:
  - defaults
  - conda-forge
dependencies:
  - python=3.9
  - pip
  - pip:
    - mlflow>=1.0
    - scikit-learn
    - pandas

Теперь, имея эти файлы и ваш скрипт (train.py), любой человек (или система CI/CD) может запустить ваш проект командой:

mlflow run . --experiment-name="tutorial" -P alpha=0.4 -P l1_ratio=0.2

MLflow автоматически создаст указанное окружение (если его еще нет), загрузит зависимости и выполнит команду из entry_points, передав указанные параметры. 

Практический пример использования MLflow на Python с проектом и Model Registry

Создадим простой проект, который обучает линейную регрессию на случайных данных, логирует все в MLflow Tracking, использует MLflow Projects для запуска и регистрирует модель в Model Registry.

Шаг 1: Структура проекта

Создайте папку, например mlflow_example, и внутри нее следующие файлы:

mlflow_example/
 ├── train.py
 ├── MLproject
 └── conda.yaml

Шаг 2: Файл conda.yaml (окружение)

name: mlflow-sklearn-example
 channels:
   - defaults
   - conda-forge
 dependencies:
   - python=3.9
   - pip
   - pip:
     - mlflow>=1.10 # Используем версию с поддержкой Model Registry
     - scikit-learn
     - numpy
     - pandas

Шаг 3: Файл MLproject (конфиг проекта)

name: sklearn-regression-example
 
 environment: conda.yaml
 
 entry_points:
   main:
     parameters:
       n_samples: {type: int, default: 100}
       noise: {type: float, default: 0.1}
       model_name: {type: str, default: "LinearRegressionDemo"} # Имя для регистрации модели
     command: "python train.py --n_samples {n_samples} --noise {noise} --model_name {model_name}"

Шаг 4: Файл train.py (скрипт обучения)

import argparse
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

import mlflow
import mlflow.sklearn

def generate_data(n_samples, noise):
    X = np.random.rand(n_samples, 1) * 10
    y = 2 * X.squeeze() + 3 + np.random.randn(n_samples) * noise * 10
    return train_test_split(X, y, test_size=0.2, random_state=42)

def train_model(n_samples: int, noise: float, model_name: str):
    """Обучает модель и логирует параметры, метрику и артефакт."""

    # Генерируем данные
    X_train, X_test, y_train, y_test = generate_data(n_samples, noise)

    # Логируем параметры
    with mlflow.start_run(nested=True): 
        mlflow.log_param("n_samples", n_samples)
        mlflow.log_param("noise", noise)
        mlflow.log_param("model_name", model_name)

        # Обучаем
        model = LinearRegression()
        model.fit(X_train, y_train)

        # Предсказываем и логируем MSE
        y_pred = model.predict(X_test)
        mse = mean_squared_error(y_test, y_pred)
        mlflow.log_metric("mse", mse)
        print(f"Mean Squared Error: {mse}")

        # Логируем модель как артефакт и регистрируем её
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="sklearn-model",
            registered_model_name=model_name
        )
        print(f"Model logged and registered under name: {model_name}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n_samples", type=int, default=100)
    parser.add_argument("--noise", type=float, default=0.1)
    parser.add_argument("--model_name", type=str, default="LinearRegressionDemo")
    args = parser.parse_args()

    train_model(
        n_samples=args.n_samples,
        noise=args.noise,
        model_name=args.model_name
    )

if __name__ == "__main__":
    main()

Шаг 5: Запуск и просмотр результатов

1. Установите MLflow: pip install mlflow (если еще не установлен).

2. Перейдите в папку mlflow_example в вашем терминале.

3. Запустите MLflow UI в фоновом режиме (или в другом окне терминала):

• mlflow ui

• Откройте в браузере http://127.0.0.1:5000.

4. Запустите проект:

• mlflow run . --experiment-name="sklearn-regression-example" -P n_samples=200 -P noise=0.2 -P model_name="LinRegNoise0.2"

• MLflow создаст conda-окружение (это может занять время при первом запуске), а затем выполнит train.py с указанными параметрами.

5. Проверьте MLflow UI:

• Вы увидите новый запуск (Run) в эксперименте sklearn-regression-example.

• Внутри запуска будут залогированы параметры (n_samples, noise, model_name_param).

• Будет залогирована метрика mse.

• В разделе Artifacts вы увидите папку sklearn-model, содержащую файлы модели.

• Перейдите в раздел Models в MLflow UI. Вы увидите зарегистрированную модель LinRegNoise0.2 с версией 1.

Попробуйте запустить проект еще раз с другими параметрами, например:

mlflow run . --experiment-name="sklearn-regression-example" -P n_samples=2000 -P noise=0.4 -P model_name="LinRegNoise2000_0.4"

Вы увидите новый Run и новую зарегистрированную модель. Если вы используете то же model_name, MLflow создаст новую версию существующей модели.

Трекинг LLM c помощью MLflow

С взрывным ростом больших языковых моделей (LLM) возникает вопрос: а можно ли использовать MLflow для них? Ответ — да! Принципы остаются теми же, но акценты смещаются.

При работе с LLM с помощью MLflow можно отслеживать:

• Параметры:

• Имя базовой модели (e.g., gpt-3.5-turbo, llama-2-7b).

• Параметры генерации (temperature, top_p, max_tokens).

• Версию промпта (или сам текст промпта, если он не слишком длинный).

• Параметры fine-tuning (если вы дообучаете модель): learning rate, batch size, количество эпох.

• Метрики:

• Стандартные метрики оценки текста (BLEU, ROUGE).

• Метрики оценки на конкретных задачах (точность классификации, F1 для извлечения сущностей).

• Оценки качества ответа, полученные с помощью других LLM (например, оценка релевантности ответа).

• Latency (время ответа модели).

• Стоимость вызова API (если используется платная модель).

• Артефакты:

• Полный текст промптов (особенно если они сложные или шаблонные).

• Примеры сгенерированных ответов (input/output pairs).

• Наборы данных для fine-tuning или оценки.

• Визуализации (если применимо).

• Логи ошибок или неудачных генераций.

• Если вы делаете fine-tuning, то саму дообученную модель (например, адаптеры LoRA).

MLflow также развивает интеграции с популярными LLM-фреймворками (например, LangChain, Transformers), чтобы сделать логирование еще проще. Основная идея — использовать знакомые mlflow.log_param, mlflow.log_metric, mlflow.log_artifact для специфичных для LLM сущностей.

Обучиться работе с моделями машинного обучения: от базовой математики до написания собственного алгоритма — можно на совместной магистратуре Skillfactory и МИФИ «Прикладной анализ данных и машинное обучение».