MLflow: учёт экспериментов, когда моделей стало много

MAATRIX

Если вы дообучаете модель хотя бы неделю подряд, у вас уже минимум десяток прогонов с разными гиперпараметрами, версиями данных и чекпоинтами — и часть из них вы не сможете сейчас отличить друг от друга. Через месяц вопрос «а как мы получили вот этот результат» превращается в археологию по именам файлов вида model_final_v2_real_fixed.pt. MLflow решает именно это: систематически фиксирует, что использовалось в каждом запуске и что из этого получилось, а не заставляет вас помнить это самому.

Проблема: почему заметок и имён файлов быстро перестаёт хватать

Пока экспериментов пять, держать их в голове или в блокноте несложно. Проблема начинается на масштабе — когда за неделю активной работы накапливается 30-50 запусков с разными комбинациями learning rate, batch size, версий датасета и архитектуры.

Типичные вопросы, на которые перестаёт хватать памяти и файловых имён:

• Какой именно набор параметров дал вот этот конкретный чекпоинт с val_accuracy 0.891?
• Какая версия модели сейчас реально стоит в проде — та, что с augmentation, или без?
• Чем текущий эксперимент отличается от предыдущего успешного варианта — поменяли одну переменную или пять сразу?
• Куда делся график с confusion matrix для запуска трёхнедельной давности?

Ручной учёт (таблица в Excel, папки experiment_1, experiment_2_fixed, заметки в Notion) какое-то время работает, но ломается ровно тогда, когда становится нужнее всего — в момент, когда экспериментов много и цена ошибки в выборе модели растёт. Таблицу забывают обновить после неудачного прогона, файлы переименовывают вручную и путают версии, а связь «эти метрики получены именно этим кодом с именно этими данными» держится только в памяти автора.

MLflow — open-source инструмент, который берёт эту рутину на себя: при каждом запуске обучения он автоматически фиксирует параметры, метрики, артефакты и связывает их в один просматриваемый и сравнимый объект — без ручного ведения таблиц.

Что именно фиксирует MLflow при каждом запуске

Центральное понятие в MLflow — run (запуск): один прогон обучения или эксперимента. Для каждого run автоматически сохраняются четыре вещи:

1. Параметры (params) — гиперпараметры, с которыми запущено обучение: learning rate, batch size, число эпох, версия датасета, архитектура, seed. Пишутся один раз при старте и не меняются.
2. Метрики (metrics) — числовые показатели качества: loss, accuracy, F1, perplexity и что угодно ещё. В отличие от параметров метрики можно логировать по шагам (по эпохам или по батчам), и MLflow сам построит график изменения метрики во времени.
3. Артефакты (artifacts) — файлы: сама обученная модель (веса), графики (confusion matrix, ROC-кривая), логи обучения, конфиги, примеры предсказаний. Хранятся отдельно от базы метаданных — на диске или в S3-совместимом хранилище.
4. Метаданные запуска — время старта и завершения, статус (FINISHED/FAILED/RUNNING), тег с версией кода (git commit), кто и с какой машины запустил.

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

Разворачиваем MLflow Tracking Server на своём сервере

Для одного человека и небольших экспериментов достаточно mlflow ui локально с файловым хранилищем. Но если над моделями работает больше одного человека или экспериментов действительно много, нужен отдельный tracking-сервер с полноценной базой данных под метаданные и объектным хранилищем под артефакты — иначе локальная файловая база (mlruns/) станет тяжёлой и её неудобно расшаривать между машинами.

Практическая схема: PostgreSQL под backend store (параметры, метрики, метаданные экспериментов) и MinIO под артефакты (S3-совместимое хранилище, поднятое на своём сервере — про его установку у нас есть отдельный разбор: MinIO в docker-compose).

Официальный образ mlflow не включает драйверы для Postgres и S3, поэтому нужен свой Dockerfile:

FROM python:3.11-slim
RUN pip install --no-cache-dir mlflow psycopg2-binary boto3

И docker-compose.yml:

services:
  postgres:
    image: postgres:16
    restart: unless-stopped
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
    volumes:
      - pgdata:/var/lib/postgresql/data

  minio:
    image: minio/minio:latest
    restart: unless-stopped
    command: server /data --console-address ":9001"
    environment:
      MINIO_ROOT_USER: ${MINIO_ROOT_USER}
      MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD}
    volumes:
      - miniodata:/data

  mlflow:
    build: .
    restart: unless-stopped
    depends_on:
      - postgres
      - minio
    environment:
      MLFLOW_S3_ENDPOINT_URL: http://minio:9000
      AWS_ACCESS_KEY_ID: ${MINIO_ROOT_USER}
      AWS_SECRET_ACCESS_KEY: ${MINIO_ROOT_PASSWORD}
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:${POSTGRES_PASSWORD}@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0 --port 5000
    ports:
      - "5000:5000"

volumes:
  pgdata:
  miniodata:

Бакет mlflow-artifacts в MinIO нужно создать заранее (через консоль MinIO на порту 9001 или mc mb). После docker compose up -d интерфейс MLflow доступен на порту 5000.

Важный нюанс: у open-source MLflow нет встроенной авторизации в UI — любой, кто достучится до порта 5000, увидит все эксперименты и сможет их менять. На боевом сервере порт наружу лучше не выставлять вообще, а доступ давать через reverse proxy с basic-auth (о том, как вообще устроен путь запроса через reverse proxy — в отдельной статье):

server {
    listen 443 ssl;
    server_name mlflow.example.internal;

    location / {
        auth_basic "MLflow";
        auth_basic_user_file /etc/nginx/.htpasswd;
        proxy_pass http://127.0.0.1:5000;
        proxy_set_header Host $host;
    }
}

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

Интеграция в код обучения

На стороне клиента нужен только пакет mlflow и несколько строк вокруг существующего тренировочного цикла — переписывать сам пайплайн обучения не требуется.

import mlflow

mlflow.set_tracking_uri("https://mlflow.example.internal")
mlflow.set_experiment("resnet-finetune")

with mlflow.start_run(run_name="lr_1e-4_batch_32_aug"):
    mlflow.log_param("learning_rate", 1e-4)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("dataset_version", "v3_cleaned")
    mlflow.log_param("augmentation", True)

    for epoch in range(epochs):
        train_loss = train_one_epoch(model, train_loader, optimizer)
        val_acc = evaluate(model, val_loader)

        mlflow.log_metric("train_loss", train_loss, step=epoch)
        mlflow.log_metric("val_accuracy", val_acc, step=epoch)

    mlflow.log_artifact("confusion_matrix.png")
    mlflow.pytorch.log_model(model, artifact_path="model")

Для популярных фреймворков (scikit-learn, PyTorch Lightning, Keras, XGBoost) в MLflow есть режим автологирования — не нужно вручную вызывать log_param/log_metric для стандартных вещей:

import mlflow.sklearn
mlflow.sklearn.autolog()

# дальше обучаете модель как обычно —
# параметры и метрики попадут в MLflow автоматически
model.fit(X_train, y_train)

Автологирование удобно как отправная точка, но для нестандартных метрик (кастомная бизнес-метрика, промежуточные графики) всё равно придётся логировать вручную — autolog покрывает типовые случаи, а не всё подряд.

Model Registry: какая версия модели сейчас в проде

Отдельный модуль MLflow — Model Registry — решает вторую часть проблемы: не «какие параметры дали этот результат», а «какая из обученных моделей сейчас реально используется и чем она отличается от других версий-кандидатов».

Модель регистрируется из конкретного run:

result = mlflow.register_model(
    model_uri="runs:/<run_id>/model",
    name="resnet-finetune"
)

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

from mlflow import MlflowClient

client = MlflowClient()
client.set_registered_model_alias("resnet-finetune", "production", result.version)

Приложение в проде обращается не к конкретному номеру версии, а к алиасу:

model = mlflow.pyfunc.load_model("models:/resnet-finetune@production")

Это снимает конкретный вопрос «а какая версия сейчас в проде» — ответ один клик в UI Model Registry, а не переписка в чате «кто последний деплоил модель». Когда появляется новый кандидат лучше текущего — алиас переставляется на новую версию, старая остаётся зарегистрированной и доступной для отката.

Практическая ценность: воспроизводимость и решения, а не просто журнал

Три конкретных сценария, где наличие такого учёта окупает время на его настройку:

Воспроизводимость задним числом. Через несколько недель или месяцев к вам возвращаются с вопросом «а как была получена вот эта модель с такими метриками». Без систематического учёта ответ — это попытка вспомнить или перебрать старые файлы. С MLflow — это открыть run и увидеть точный набор параметров, версию датасета и артефакты, которые к этому результату привели.

Прозрачность продакшен-версии. В команде из нескольких человек рано или поздно возникает рассинхрон: кто-то думает, что в проде модель А, а на самом деле там уже стоит модель Б после чужого деплоя. Model Registry с алиасами делает эту информацию видимой для всех, а не устной договорённостью.

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

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

Если вы дообучаете модель через LoRA, стоит сразу логировать в MLflow и это — какой базовый чекпоинт брался, какой ранг адаптера и какой датасет; подробнее про сам процесс — в статье про LoRA-файнтюнинг своей модели на VPS.

Частые вопросы