Время прочтения: 4 мин.

Есть множество известных инструментов для поиска гиперпараметров, например, Optuna, Hyperopt. У этих инструментов есть один минус — они решают только проблему оптимизации, не давая инструментов MLOps, вроде сохранения моделей, метрик и датасетов. Но есть один инструмент, который может выступать оболочкой для обеих библиотек, приведённых выше, который даёт полный функционал для логирования и экспериментирования с моделями. Это Wandb Sweeps.

Сначала о Wandb — это набор решений, которые позволяют сохранять информацию о модели, такую как гиперпареметры, метрики и веса. Он имеет интеграции с популярными ml-фреймворками (Pytorch, Pytorch Lightning, Keras, TensorFlow, HuggingFace, XGBoost, Scikit-learn), что позволяет добавить логирование в несколько строк. Также может быть встроено в любой скрипт на Python.

Wandb Sweeps позволяет оптимизировать гиперпараметры модели одним из способов: grid (перебор всех значений по сетке), random (выбор случайного набора параметров), bayes (байесовский поиск). С первыми двумя всё очевидно, на последнем остановимся подробнее.

Так упрощенно выглядит процедура байесовского поиска:

  1. Создать список гиперпараметров — каждый из них должен обладать определенным распределением, включая, например, категориальное, т.е. можно просто ввести конечное количество значений
  2. Для каждого параметра построить распределение вероятностей его оптимальности, т.е. провести несколько запусков и построить вероятностную модель целевой функции (оптимизируемой метрики)
  3. Построить распределение вероятностей того, что запуск с такими параметрами даст определенный скор — чем он лучше, тем это менее вероятно. Для шагов 2 и 3 байесовскому поиску нужно указать метрику, которую будем оптимизировать, для random и grid это не нужно
  4. Объединить распределения в полное распределение вероятностей (обычно это подразумевает их простое перемножение). В результате получаем вероятность получения определённого скора для всех возможных вариантов гиперпараметров
  5. Сделать шаг по пространству гиперпараметров в сторону наиболее вероятного для высокого скора значения
  6. Делать оптимизационные шаги и обновлять оценки вероятностей по теореме Байеса. То есть, если при lr=0.01 результат получился слабым, то шанс получить сильный результат при lr=0.099 снижается и модель проверит это значение с меньшим шансом

Ещё одно выгодное преимущество random и bayes над grid — возможность задавать значения как распределение. Например, я уверен, что lr=0.01 это хороший вариант, а лучший лежит где-то рядом. Тогда можно задать значения для lr из нормального распределения с средним 0.01 и среднеквадратичным отклонением, например, 0.005. Так это выглядит в конфиге из которого запускается поиск:

lr:
  distribution: normal
  mu: 0.01
  sigma: 0.005

Больше примеров в документации, ссылка.

Короткий пример практического применения случайного поиска:

import wandb

def train():
   with wandb.init() as run:
       config = wandb.config
       model = make_model(config)
       for epoch in range(config["epochs"]):
           loss = model.fit()  # your model training code here
           wandb.log({"loss": loss, "epoch": epoch})

sweep_config = {
 "name" : "my-sweep",
 "method" : "random",
 "parameters" : {
   "epochs" : {
     "values" : [10, 20, 50]
   },
   "learning_rate" :{
     "min": 0.0001,
     "max": 0.1
   }
 }
}

sweep_id = wandb.sweep(sweep_config)
count = 5 # number of runs to execute
wandb.agent(sweep_id, function=train, count=count)

Вот полученный результат, ссылка.

Более подробный результат применения сети к проекту, не влезающий в рамки статьи – ссылка на коммит в гитхабе, где я добавляю wandb к уже существующему проекту, yolact, нейросети для сегментации изображений. Вот результат нескольких прогонов, ссылка.

Но нужно поговорить и о недостатках.

Альтернативой используемому инструменту можно считать связку, например, Optuna + MLFlow, в который первая часть отвечает за поиск гиперпараметров, а вторая за логирование. У этой альтернативы есть несколько преимуществ – можно задавать любое распределение для параметра (получив его через numpy), бесплатное использование для команд разработки, более простое допиливание под конкретную задачу, наличие генетических алгоритмов и других стратегий оптимизации, которых нет в Wandb.  Эти недостатки не играют никакой роли, если вы не хотите сильно вникать в подбор параметров и просто хотите сделать это быстро и удобно. Обычной байесовской оптимизации с нормальным или равномерным распределением с лихвой хватает для почти любой задачи.

При байесовском поиске в Wandb, как и в Optuna, можно посмотреть parameter importance для каждой сохранённой метрики, иллюстрации этого можно найти по ссылке. Кроме того, важным преимуществом Wandb Sweeps над связкой, например, той же Optuna + MLFlow является возможность запускать буквально одной (!) строкой в консоли поиск по одному пространству гиперпараметров на нескольких машинах одновременно.

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