Менеджер ML-экспериментов. Что это и как он нас выручает

Меня зовут Владимир Кочетков, и я тимлид Deep Learning разработки в области распознавания речи и соавтор курса «MLOps для разработки и мониторинга моделей» ^[1] в Яндекс Практикуме. А ещё автор скромного образовательного телеграм-канала ^[2] про ML :-)

Сегодня поделюсь опытом ^[3], как мы с командой внедряли менеджер ML-экспериментов, и расскажу:

• в чём сложность обучения ^[4] моделей;

• когда нам понадобился менеджер экспериментов;

• какие были к нему требования и как они эволюционировали в процессе;

• что мы выбрали и почему;

• как это работает в жизни.

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

Обучение моделей — штука очень непредсказуемая

В сети каждый день мы встречаем фразы типа «Мы обучили такую-то модель делать вот это» — и уже привыкли к этому. Но давайте разберёмся, что за этим стоит. 

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

• архитектуру: Transformer, Conformer, CNN, RNN и т.д.;

• параметры модели: количество слоев, голов, функции активации;

• гиперпараметры оптимизации: learning rate, batch size;

• подходы к обучению: комбинация функций потерь, дистилляция, RL;

• способы предобработки данных: алгоритмы фильтрации, аугментации;

• версии данных: источник 1, источник 2, объединение обоих.

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

От такого изобилия различных вариантов часто опускаются руки, потому что не знаешь, что лучше использовать (и какой именно параметр всё снова заруинил). Те, кто только начинает свой путь в этой теме, иногда вывозят на энтузиазме, но в реальной разработке такой опции нет — весь энтузиазм закончился на десятом подходе, и нового пока не завезли.

И что остаётся делать? Опираясь на свой и чужой опыт, следуя рекомендациям (а также интуиции ^[5], ей тоже нужно доверять), мы проводим эксперименты: 

1. Готовим данные.

2. Обучаем

3. Тестируем.

4. Сравниваем результаты. 

Бывало и такое, что провёл какой-то эксперимент, потом забыл про это и провёл тот же самый эксперимент. Получается, тратим и время, и вычислительные мощности.

Чтобы не запутаться во всех этих попытках, их настройках, нам необходимо хранить достаточно информации о каждом эксперименте. На этом остановимся подробнее.

Что значит «Хранить достаточно информации об эксперименте»?

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

На этом моменте у многих возникает вопрос: «Зачем нам получать те же самые результаты, если они у нас уже есть?». А затем, что при обучении нейросетей нам часто требуется воспроизвести эксперимент при некоторых новых вводных, но при остальных прежних условиях. 

Например, у нас имеется эксперимент обучения классификатора:

• модель: transformer;

• гиперпараметры оптимизации: learning rate = 0.00001, batch size = 64;

• версия данных: 1.4;

• версия кода: 0.7.2

• метрики: precision = 0.973, recall = 0.924

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

Более того, сравнив метрики этих двух экспериментов, мы можем сделать вывод о влиянии на результат именно версии данных, поскольку остальные настройки гарантированно одинаковые.

А теперь давайте представим, что будет, если информации будет недостаточно. Пусть мы имеем следующий эксперимент

• модель: conformer;

• версия данных: 1.4;

• версия кода: 0.7.2;

• метрики: precision = 0.981, recall = 0.955.

Судя по метрикам, этот эксперимент с новой моделью покажет результаты ещё лучшие, чем первый. Однако, чтобы его воспроизвести на новой версии данных, нам не хватает значений параметров learning rate и batch size, которые привели к лучшим метрикам. А значит, новый эксперимент может не дать тех результатов, на которые мы рассчитываем. 

Тогда нам остаётся выбрать эти значения произвольным образом, а это рискованно: можно установить неудачные значения и получить метрики хуже. Ещё можно провести несколько экспериментов с разными значениями и выбрать лучшие на основе метрик, но это требует больше времени (и далеко не всегда оно есть).

Что бы мы тут ни делали, у нас нет оснований делать вывод о влиянии на результат именно версии данных, поскольку нет гарантий, что learning rate и batch size были одинаковы в обоих экспериментах.

Получается, что хранение недостаточной информации чревато либо меньшим шансом на успех, либо тратой ресурсов.

Первые попытки контроля

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

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

Выглядело это как-то так: 

После довольно непродолжительного времени мы перешли на Excel-таблицы — и это был прогресс! Они позволяли выделять перспективные эксперименты зеленым цветом и даже сортировать строки по столбцам метрик, в отличие от записей в тетрадке. 

Здесь уже можно добавить в код внесение данных, которые затем открыть в Excel, но нормальным трекером экспериментов это не назовешь (хотя бы потому, что сравнивать определённые результаты между собой всё ещё неудобно). Но это была наша первая попытка полноценной автоматизации, и с неё начался наш поиск полноценного трекера.

Пробуем трекер экспериментов

Очевидно, что идея хранить данные об экспериментах пришла не только нам. После поиска и сравнения с аналогами достаточно долгое время мы использовали TensorBoard ^[6], так как он вышел раньше остальных, был хорошо обкатан и легко встраивался в код. Тогда мы работали на фреймворке tensorflow, и TensorBoard туда вставал как родной:

Для записи данных в трекер использовался простой код, что-то вроде такого:

test_summary_writer = tf.summary.create_file_writer('test/logdir')

with test_summary_writer.as_default():
    tf.summary.scalar('loss', 0.345, step=1)
    tf.summary.scalar('loss', 0.234, step=2)
    tf.summary.scalar('loss', 0.123, step=3)
 [7]

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

Да, сейчас такой функционал — база трекера ML-экспериментов, но тогда для нас это было в новинку, потому что не нужно было ничего переносить руками, а это уже достижение :-)

Но нам всё ещё приходилось где-то хранить дополнительную информацию, которую не предусматривала логика ^[8] TensorBoard, но которую нужно было иметь для полной картины. Это, например, аргументы командной строки, текущее состояние кода (коммит), версия датасета и другие настройки, определяющие исход эксперимента.

Через пару лет нам захотелось попробовать что-то новое и более полноценное, поэтому развернули и поигрались с VertaAI ^[9]. Эта платформа понравилась тем, что позволяла хранить не один проект, а сразу несколько. Там был тот же функционал, что и в TensorBoard, но можно было хранить больше метаинформации. Буквально через пару дней выяснилось, что отображение графиков, матриц ошибок (confusion matrix) и таблиц работает только для enterprise-версии. Об этом я узнал из своего issue ^[10] на GitHub:

Пользовались программой мы тоже недолго: решили обновить платформу, чтобы получить новый функционал → неудачное обновление привело к тому, что всё перестало работать → восстановить данные тоже не получилось. К слову сказать, последний коммит в репозитории ^[9]VertaAI был 2 года назад.

После этого мы с командой решили составить список требований и пожеланий к трекеру ML-экспериментов, чтобы он удовлетворял всех.

Выбираем лучший ML-менеджер для своих задач

Вот приблизительный список критериев, которые мы для себя определили:

1. Единый сервис, в котором будут собраны все наши проекты.

2. Базовые функции трекера (как было в TensorBoard): логирование метрик, сохранение гиперпараметров, построение графиков, сравнение результатов.

3. Дополнительное хранение всей достаточной информации эксперимента: скрипт запуска, аргументы командной строки, состояние кода, перечень библиотек или версия docker-образа, описание, версия данных, тэги.

4. Хранение артефактов, например моделей.

5. Добавление медиаданных для отладки: примеры картинок или аудио.

Всё это относится именно к трекеру ML-экспериментов. Но мы решили пойти дальше и рассмотреть менеджер экспериментов. Он имеет тот же функционал, что и трекер, добавляя при этом возможности:

• настройки экспериментов (изменение гиперпараметров, аргументов командной строки);

• управления рабочими процессами (воркеры, очереди) для автоматизированного запуска;

• воспроизведения экспериментов, в том числе с изменёнными в менеджере параметрами;

• интеграции с другими ML-компонентами, например системой управления данными.

В итоге остановились на трёх ML-платформах: Weights & Biases, ClearML, MLFlow.

Weights & Biases ^[11] (W&B) — мощный инструмент для отслеживания, визуализации и управления жизненным циклом ML. Но на момент выбора инструмента у W&B был один самый главный недостаток, который не позволил нам его выбрать, — его нельзя было развернуть в нашей инфраструктуре, а облачные решения мы не рассматривали. Сейчас, насколько известно, такая возможность появилась, но W&B свой шанс упустили :-)

ClearML ^[12] удовлетворял нас по всем критериям. Более того, это целая ML-платформа, которая помимо трекинга позволяла ещё организовать версионирование датасетов (о чём пойдет речь далее) и создание пайплайнов. 

В нём можно было хранить действительно всю достаточную информацию об эксперименте. Поток вывода (stdout) тоже сохранялся автоматически, что очень полезно при отладке. Даже состояние кода фиксировалось не только в рамках коммита — незакоммиченные изменения тоже сохранялись. Он стал для нас фаворитом — ровно до того момента, когда начали разворачивать его у себя.

И хотя документация заявляла поддержку self-hosted, в этой битве за локальный софт мы тогда проиграли. Нам оставалось попробовать MLFlow ^[13] — и либо с ним найти идеал, либо продолжать поиски лучшего инструмента.

К тому моменту я уже разобрался, как в MLFlow организовать трекинг экспериментов, и даже успел провести обучение для своей команды. По сравнению с ClearML это был относительный компромисс, но зато оно работало!

И только мы настроились вести всё в MLFlow, нам опытные коллеги подсказали лайфхак, как установить ClearML. И да, на этот раз всё получилось: кажется, нужно было для установки выбирать более ранние версии, чем мы тогда взяли (вот инструкция ^[14], как всё сделать с первого раза).

Интегрировать логирование через ClearML вообще не составило труда — огромная доля информации сохранялась автоматически: состояние кода, аргумента командной строки, используемый фреймворк, мониторинг вычислительных ресурсов и поток вывода (stdout). 

Пример логирования в коде (оцените простоту интеграции):

from clearml import Task, Logger

Task.set_credentials(
   key='xxxxxx',
   secret='yyyyyy'
)
task = Task.init(project_name='Car Classification',
                  task_name='train-experiment-001',                
                  task_type=TaskTypes.training,
                  auto_connect_frameworks=False)
logger = Logger.current_logger()

task.add_tags(['example'])
hparams = {
   'learning_rate': 0.001
}
task.connect(hparams, name='hparams')

for epoch in range(10):
    # Ваши вычисления
    train_loss = np.random.rand() * epoch
    val_accuracy = np.random.rand()
    
    # Логируем скалярные метрики
    logger.report_scalar(
        title='Loss',          # Группа метрик
        series='train',        # Название серии (линии на графике)
        value=train_loss,
        iteration=epoch        # Шаг (эпоха)
    )
    
    logger.report_scalar(
        title='Accuracy',
        series='validation',
        value=val_accuracy,
        iteration=epoch
    )
 [7]

Впечатления от работы в ClearML

Мне не терпелось попробовать запустить обучение средствами ClearML. Я хотел зайти на платформу через веб-браузер, взять успешный эксперимент, заменить в нём какие-то параметры (например, количество эпох), перезапустить и посмотреть полученный результат. 

Мне кажется, что это приятнее, чем зайти в IDE, поменять конфигурации в нём, запустить docker-контейнер. А ещё это надежнее — вы гарантированно получаете все конфигурации и результаты по всем экспериментам, которые проходят через менеджер.

Представьте, что вы провели обучение, скажем, пару недель назад. После этого ни один эксперимент не дал результатов лучше, хотя вы существенно поменяли код и/или конфигурации. А если вы захотите воспроизвести тот старый эксперимент, то придётся вернуть прежнее состояние кода и конфигураций, что не только муторно само по себе, но и увеличивает риск ошибки ^[15]. Раз уж ClearML сохраняет всю эту информацию, почему бы нам просто не продублировать этот эксперимент, заменив нужные параметры, и перезапустить его с новыми. 

Если один раз должным образом настроить всё необходимое — экономия времени и детальность записей будут колоссальные. Тут, конечно, пришлось понаступать на грабли (здесь ^[16] можно почитать про это более подробно), но результат точно того стоит.

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

Для тех, кому хочется копнуть поглубже: держите более подробное использование менеджера экспериментов ^[17] в шаблонном проекте ^[18].

Не забываем про организацию датасетов

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

• Различные источники (kaggle, audioset, common voice).

• Множество способов их предобработки (чистка, фильтрация, преобразования).

• Добавление аугментаций.

• И ещё множество комбинаций объединения всего вышесказанного.

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

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

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

Важный аспект организации датасетов — это хранение истории создания каждой версии данных. Например, была версия 1, к которой применили какое-то преобразование и получили версию 2. По результатам экспериментов метрики на версии 2 хуже, чем на версии 1. Зная, как была получена версия 2, мы делаем вывод, что данное преобразование данных плохо сказывается на метриках. 

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

Есть различные инструменты для версионирования данных, среди которых:

• DVC (git для данных);

• LakeFS (та же концепция, но на уровне метаданных поверх S3);

• ClearML Datasets.

Сначала мы рассматривали LakeFS за счёт удобства использования знакомых git-команд и заявленных интеграций с Airflow, Kafka и S3-совместимым хранилищем. Но всё же решили остановиться на ClearML Datasets: 

• эта платформа уже была развернута в нашем контуре;

• удобство взаимодействия между экспериментами и датасетами в рамках одной платформы;

• приятное отображение всей истории создания датасета в виде ветвлений;

• сохранение всей необходимой информации о создании версии. Помимо того, что метаданные логируются руками, автоматически сохраняются: скрипт запуска, состояние кода, аргументы командной строки и поток вывода (stdout).

В коде это может выглядеть так:

from clearml import Dataset

# Создаём датасет (указываем проект и имя)
dataset = Dataset.create(
    dataset_name="my-image-dataset",
    dataset_project="my-project",
    parent_datasets=None,  # Можно указать родительские датасеты (опционально)
    tags=["vision", "classification"]  # Теги для поиска
)

print(f"Создан датасет: {dataset.id}")

# Добавляем локальные файлы/директории
dataset.add_files(
    path="/path/to/images",           # Локальный путь
    dataset_path="raw/images",       # Путь внутри датасета
    verbose=True
)

# Можно добавить отдельные файлы
dataset.add_files(
    path="/path/to/labels.csv",
    dataset_path="raw/labels.csv"
)

# Или файлы по URL
dataset.add_url(
    url="https://example.com/data.zip",
    dataset_path="external/data.zip"
)

# Финализируем и создаём версию
dataset.upload(
    output_url=None,  # Если None — используется дефолтное хранилище
    verbose=True
)

# Публикуем версию (обязательно!)
dataset.finalize()

print(f"Датасет опубликован. Версия: {dataset.version}")
 [7]

С использованием версионирования процессы работы с данными становятся более прозрачными и становится проще проводить связь между экспериментами и данными. Ещё минимизируется дублирование данных, поскольку одна платформа применяется для всех.

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

Что мы получили в итоге

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

Напоследок напомним тернистый путь поисков :-)

С чего начинали:

И к чему пришли в итоге:

Бонус-трек: с чего можно начать внедрение менеджера

Простой и пошаговый план для тех, кто хочет это попробовать:

1. Внедрить ML-платформу, включающую как минимум функционал трекинга экспериментов. 

2. Модифицировать код обучения и тестирования, внедрив логирование метрик, метаданных и другой важной информации, постепенно обеспечивая воспроизводимость экспериментов. 

3. Попробовать механизмы управления рабочими процессами в ML-платформе: они  позволят ставить их «пачками» в очередь и не отвлекаться на ручной запуск. 

4. Организовать датасеты, методы их получения и связать их версии с экспериментами. 

5. Увлечённым и энтузиастам рекомендую обратить внимание ^[19] на реестр ML-моделей, который позволяет хранить полученные в ходе экспериментов модели и их метаданные.