Основы фреймворка Hydra

Владислав Ефимов — ML Engineer

Начальные сведения

Hydra – инструмент для создания и компоновки файлов конфигурации различных Python-приложений с акцентом на использование в ML-проектах.

Хорошей практикой при разработке любого более-менее сложного приложения считается выделение файлов с ключевыми параметрами, существенно влияющими на работу кода, в отдельные конфигурацинные файлы. Она облегчает восприятие структуры кода самим автором и другими разработчиками, упрощает тестирование приложения в различных сценариях его работы. Примерами таких параметров могут быть: номер сетевого порта, url-адреса, имя пользователя и его роль в БД и многое другое. ML-проекты в этом смысле ничуть не менее требовательны, ведь отдельные конфигурации могут потребоваться на многих этапах пайплайна: на этапе получения и хранения данных, на этапе препроцессинга, при определении и обучении моделей. Более того, каждый из этих этапов может быть запущен большое число раз с различными параметрами в рамках экспериментов по построению оптимальной ML-модели.

Выделение параметров в группы, компоновка групп в общую иерархическую структуру
Переопределение групп конфигураций и отдельных параметров
Автоматическое логирование
Запуски с перебором различных комбинаций параметров

Hydra дает возможность легко манипулировать различными наборами конфигураций:

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

Hydra по сути является надстройкой над библиотекой OmegaConf, которая объединяет конфигурации из разных источников, унифицирует их в YAML-формат и осуществляет проверку типов, поэтому часть API OmegaConf’a доступна и в ней.

Практическое применение

Задача: Sentiment Analysis (анализ тональности текста)
Датасет: Amazon reviews (отзывы клиентов Amazon на различные товары)
Модели: DistillBERT, RoBERTa
Hugging Face transformers API: Trainer, TrainingArguments, Auto Model

Будем исследовать возможности Hydra на примере задачи классификации отзывов на товары компании Amazon: требуется по тексту самого отзыва и его заголовка (англ.) определить отзыв как положительный, нейтральный или негативный. Воспользуемся двумя разновидностями архитектуры BERT: DistillBERT (облегченная версия исходной модели) и RoBERTa (более тяжеловесная модификация исходной модели). Фреймворк huggingface предоставляет единообразный удобный API для обучения и инференса этих моделей, приведенные ниже файлы конфигурации ориентированы на использование 3 основных классов (слайд), с документацией можно ознакомиться по ссылке.

Базовая структура

├── src
    ├── conf
      ├── config.yaml
      └── data
      └── model
          ├── distillbert.yaml
          ├── roberta.yaml
      └── preprocessing
          ├── distillbert.yaml
          ├── roberta.yaml
      └── secrets
          ├── example.yaml
          ├── .gitignore
      └── training
          ├── distillbert.yaml
          ├── roberta.yaml
├── SA_example.py

Для начала нужно создать папку, содержащую файлы конфигураций. Название и локация этой папки внутри проекта могут быть произвольными, в нашем случае это будет src/conf. Внутри неё каждая подпапка будет соответствовать отдельному направлению, или группе конфигураций, связанных с тем или иным аспектом проекта. Например, добавим папки data, model, preprocessing, secrets и training для групп параметров, связанных с источниками данных, определением модели, предобработкой данных, параметрами авторизации и обучением модели соответственно. Папка secrets будет хранить реквизиты для различных облачных сервисов, хранилищ данных и прочего. Поскольку для компоновки конфига при помощи Hydra нет необходимости явно перечислять множество параметров, достаточно добавать secrets в локальный .gitignore, чтобы получить удобный способ работать с конфиденциальными данными и не публиковать их в удаленном репозитории. Внутри каждой группы (папки) можно поместить произвольное число .yaml-файлов, соответствующих различным вариантам настроек. Добавим набор параметров для простой предобработки текстовых данных (preprocessing/distillbert.yaml, preprocessing/roberta.yaml), параметры для инициализации моделей (model/distillbert.yaml, model/roberta.yaml), а также параметры обучения модели (training/distillbert.yaml, training/roberta.yaml), так что итоговое дерево файлов будет выглядеть следующим образом:

Главный файл конфигурации

conf/config.yaml

defaults:
  - preprocessing: distillbert
  - training: distillbert
  - model: distillbert
  - secrets: example

Приоритет параметров

conf/preprocessing/roberta.yaml

padding: max_length
tokenizer_name: siebert/sentiment-roberta-large-english
test_size: 5000
train_size: 20000
max_length: 256

conf/config.yaml

defaults:
  - preprocessing: distillbert
  - training: distillbert
  - model: distillbert
  - secrets: example
  - _self_

preprocessing:
  test_size: 10000
  train_size: 50000
  max_length: 512

$ python SA_example.py

[2024-03-21 23:52:30,126][__main__][INFO] - preprocessing:
  padding: max_length
  max_length: 512
  tokenizer_name: siebert/sentiment-roberta-large-english
  test_size: 10000
  train_size: 50000

conf/preprocessing/roberta.yaml

padding: max_length
tokenizer_name: siebert/sentiment-roberta-large-english
test_size: 5000
train_size: 20000
max_length: 256

conf/config.yaml

defaults:
  - _self_
  - preprocessing: distillbert
  - training: distillbert
  - model: distillbert
  - secrets: example

preprocessing:
  test_size: 10000
  train_size: 50000
  max_length: 512

$ python SA_example.py

[2024-03-21 23:52:30,126][__main__][INFO] - preprocessing:
  test_size: 5000
  train_size: 20000
  max_length: 256
  padding: max_length
  tokenizer_name: siebert/sentiment-roberta-large-english

В главном конфиге можно указывать не только дефолтные группы, но и отдельные параметры ровно так же, как и в дочерних конфигах. В случае, когда имя параметра в главном конфиге совпадает с именем параметра в дочернем, результирующее значение параметра определяется по местонахождению поля _self_:

если оно указано в конце списка defaults, то будет использовано значение из src/conf
если оно указано в начале списка defaults, то будет использовано значение из src/conf/*/*.yaml

Модель RoBERTa обладает бОльшим числом параметров, поэтому (в учебных целях) разумно уменьшить для неё максимальную длину входных последовательностей (в токенах), а также размеры тестовой и обучающей выборок, ускорив таким образом время обучения и инференса. Заметим, что в conf/preprocessing/roberta.yaml указаны отличные от главного конфига значения параметров max_length, test_size, train_size.

В примере слева параметр max_length определен как в главном конфиге (512), так и в conf/preprocessing/roberta.yaml (256), а поскольку _self_ указано в конце списка дефолтных параметров, будет использовано значение 512.

В примере справа итоговым значением будет выбрано 256.

То же относительно остальных двух параметров.

Интеграция Hydra в Python-код, способы запуска

hydra_main(), запуск .py-модулей из терминала

@hydra.main(version_base=None, config_path="src/conf", config_name="config")
def main(cfg: DictConfig):
    param = cfg["group_name"]["param"]
    # some logic with param...

Увеличение числа эпох обучения модели DistillBERT (по умолчанию):

$ python SA_example.py training.num_train_epochs=10

Пример замены группы параметров (переход к модели RoBERTa):

$ python SA_example.py model=roberta preprocessing=roberta training=roberta

Существующие в конфиге параметры можно удалять (learning_rate будет сброшен на стандартное значение):

$ python SA_example.py ~training.learning_rate

Добавление нового параметра, не присутствующего ранее:

$ python SA_example.py +training.use_cpu=True

Также в таком варианте использования Hydra доступна возможность multirun запусков - выполнения кода со всеми перечисленными значениями параметров:

$ python SA_example.py -m preprocessing.max_length=128,512

Расширения для multirun-режима

Joblib (распараллеливание)
Ray (запуск на кластере)
Optuna (оптимизация гиперпараметров)

Compose-API, использование в jupyter-notebooks

В случае:

Jupyter-like форматов
Юнит-тестов
Когда нет доступа к терминалу

можно воспользоваться комбинацией методов compose() и initialize()

src/config.py

def compose_config(
    overrides: list[str] | None = None,
    config_path: Path = Path("conf"),
    config_name: str = "config",
) -> dict:
    with initialize(version_base=None, config_path=str(config_path)):
        hydra_config = compose(config_name=config_name, overrides=overrides)
        return hydra_config

example.ipynb

from src.config import compose_config

main_conf = compose_config()

prepr_conf = main_conf["preprocessing"]
train_conf = main_conf["training"]
model_conf = main_conf["model"]

Пример изменения группы параметров:

from src.config import compose_config
main_conf = compose_config(overrides=["model='some_other_model'"])
model_conf = main_conf["model"]

Данный способ не позволяет использовать multirun-запуски, настраивать логирование, менять стандартные пути для побочных артефактов работы Hydra и некоторые другие возможности.

Некоторые полезные расширения синтаксиса для yaml-файлов

Variable interpolation

OmegaConf docs

Абсолютный путь: ${config_name.param_name}
Относительный путь: ${.param_name}

_target_: transformers.TrainingArguments
output_dir: distillbert_training_outputs
evaluation_strategy: epoch
save_strategy: epoch
learning_rate: 2e-5
per_device_train_batch_size: 64
per_device_eval_batch_size: ${training.per_device_train_batch_size}
num_train_epochs: 3
weight_decay: 0.01
load_best_model_at_end: True
metric_for_best_model: f1
fp16: True

В любом конфиг-файле вместо прямого указания значений параметров можно сослаться на значение другого параметра из этого же или другого конфига. Такая возможность предоставляется механизмом интерполяции переменных (variable interpolation) библиотеки OmegaConf. Им можно воспользоваться при помощи синтаксиса ${config_name.param_name} (вариант с абсолютным путем), или ${.param_name} (вариан с относительным путем). Пример использования этого механизма приведен в конфиге модели conf/training/distillbert.yaml.

Resolvers

secrets/example.yaml

name: ${oc.env:USER}

$ python SA_example.py

[2024-03-22 00:12:14,867][__main__][INFO] - secrets:
  user: iref

Инициализация объектов

from hydra.utils import instantiate
model = instantiate(confifg['model'])

_target_: transformers.AutoModelForSequenceClassification.from_pretrained
_convert_: object
pretrained_model_name_or_path: lxyuan/distilbert-base-multilingual-cased-sentiments-student
problem_type: single_label_classification
num_labels: 2
ignore_mismatched_sizes: True

При создании файла конфигурации можно воспользоваться полем _target_, указав его значением экземпляр класса или вызываемый объект (callable). При этом остальные поля, указанные на том же уровне, будут использованы как параметры при инициализации класса. Отдельно импортировать этот класс в самом python-модуле не нужно, достаточно воспользоваться методом instantiate(). В некоторых случаях может понадобиться также дополнительное поле _convert_, определяющее, как instantiate будет обращаться с параметрами класса, не указанными в конфиге, а передаваемыми в runtime: конвертировать их в OmegaConf-контейнер полностью (None), частично (partial), или же вообще не конвертировать (object).

Логирование

Каждый запуск сопровождается сохранением некоторого набора логов. В них содержатся:

config.yaml - скомпонованный конфиг, с которым прошел запуск
hydra.yaml - некоторые внутренние настройки Hydra, с которыми прошел запуск
overrides.yaml - список параметров и их значений, измененных при помощи командной строки или функции compose
<filename.log> - данные, переданные через модуль logging.

$ python rubert_tiny_hydra_example.py hydra.run.dir=some_other_out_dir

config.yaml

hydra:
  run:
    dir: ./some_other_out_dir/${now:%Y-%m-%d}/${now:%H-%M-%S}

$ python rubert_tiny_hydra_example.py hydra.verbose=True

Итог

Прежде всего, Hydra — удобный инструмент для организации множества файлов с параметрами в единую конфигурацию, который упрощает процесс разработки и исследований.