Трекеры экспериментов

Существует множество API для отслеживания экспериментов, однако заставить их все работать в среде многопроцессорности зачастую бывает сложно. Accelerate предоставляет универсальный API для отслеживания, который можно использовать для записи полезных данных во время выполнения вашего скрипта с помощью Accelerator.log()

Интегрированные трекеры

Сейчас Accelerate поддерживает восемь трекеров из коробки:

• TensorBoard
• WandB
• Trackio
• CometML
• Цель
• MLFlow
• ClearML
• DVCLive

Чтобы использовать любой из них, передайте выбранный тип(ы) в параметр log_with в Accelerate:

from accelerate import Accelerator
from accelerate.utils import LoggerType

accelerator = Accelerator(log_with="all")  # For all available trackers in the environment
accelerator = Accelerator(log_with="wandb")
accelerator = Accelerator(log_with=["wandb", LoggerType.TENSORBOARD])

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

hps = {"num_iterations": 5, "learning_rate": 1e-2}
accelerator.init_trackers("my_project", config=hps)

Когда вы готовы записывать данные, следует использовать Accelerator.log(). Также можно передать step, чтобы связать данные с конкретным шагом в цикле обучения.

accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)

Когда обучение завершено, обязательно вызовите Accelerator.end_training(), чтобы все трекеры могли выполнить свои завершающие функции, если они есть.

accelerator.end_training()

Полный пример приведён ниже:

from accelerate import Accelerator

accelerator = Accelerator(log_with="all")
config = {
    "num_iterations": 5,
    "learning_rate": 1e-2,
    "loss_function": str(my_loss_function),
}

accelerator.init_trackers("example_project", config=config)

my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)
device = accelerator.device
my_model.to(device)

for iteration in range(config["num_iterations"]):
    for step, batch in enumerate(my_training_dataloader):
        my_optimizer.zero_grad()
        inputs, targets = batch
        inputs = inputs.to(device)
        targets = targets.to(device)
        outputs = my_model(inputs)
        loss = my_loss_function(outputs, targets)
        accelerator.backward(loss)
        my_optimizer.step()
        accelerator.log({"training_loss": loss}, step=step)
accelerator.end_training()

Если трекеру требуется каталог для сохранения данных, например TensorBoard, передайте путь к каталогу в параметр project_dir. Параметр project_dir полезен, когда необходимо объединить другие конфигурации в классе данных ProjectConfiguration. Например, вы можете сохранять данные TensorBoard в project_dir, а всё остальное — в параметре logging_dir класса [~utils.ProjectConfiguration:

accelerator = Accelerator(log_with="tensorboard", project_dir=".")

# use with ProjectConfiguration
config = ProjectConfiguration(project_dir=".", logging_dir="another/directory")
accelerator = Accelerator(log_with="tensorboard", project_config=config)

Реализация пользовательских трекеров

Чтобы реализовать новый трекер для использования в Accelerator, можно создать его, реализовав класс GeneralTracker. Каждый трекер должен реализовывать три функции и иметь три свойства:

• __init__:

  • Должен сохранять run_name и инициализировать API трекера интегрированной библиотеки.
  • Если трекер хранит данные локально (например, TensorBoard), можно добавить параметр logging_dir.

• store_init_configuration:

  • Должен принимать словарь values и сохранять его как одноразовую конфигурацию эксперимента

• log:

  • Должен принимать словарь values и step и записывать их в run

• name (str):

  • Уникальное строковое имя трекера, например "wandb" для трекера wandb.
  • Это будет использоваться для взаимодействия именно с этим трекером

• requires_logging_directory (bool):

  • Требуется ли для данного конкретного трекера logging_dir и использует ли он его.

• tracker:

  • Это следует реализовать как функцию @property
  • Должен возвращать внутренний механизм отслеживания, который использует библиотека, например, объект run для wandb.

Каждый метод также должен использовать класс state.PartialState, если, например, логгер должен выполняться только на главном процессе.

Ниже приведён краткий пример интеграции с Weights and Biases, содержащий только релевантную информацию и выполняющий логирование только на главном процессе:

from accelerate.tracking import GeneralTracker, on_main_process
from typing import Optional

import wandb


class MyCustomTracker(GeneralTracker):
    name = "wandb"
    requires_logging_directory = False

    @on_main_process
    def __init__(self, run_name: str):
        self.run_name = run_name
        run = wandb.init(self.run_name)

    @property
    def tracker(self):
        return self.run.run

    @on_main_process
    def store_init_configuration(self, values: dict):
        wandb.config(values)

    @on_main_process
    def log(self, values: dict, step: Optional[int] = None):
        wandb.log(values, step=step)

Когда вы будете готовы создать объект Accelerator, передайте экземпляр вашего трекера в Accelerator.log_with, чтобы он автоматически использовался в API:

tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=tracker)

Их также можно комбинировать с существующими трекерами, в том числе с "all":

tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=[tracker, "all"])

Регистрация пользовательского трекера по имени

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

from accelerate import Accelerator
from accelerate.tracking import GeneralTracker, register_tracker_class


class MyCustomTracker(GeneralTracker):
    name = "my_tracker"
    requires_logging_directory = False
    # ... implement the rest of the `GeneralTracker` interface


register_tracker_class(MyCustomTracker)
accelerator = Accelerator(log_with="my_tracker")

Доступ к внутреннему трекеру

Если требуется прямой доступ к некоторым пользовательским взаимодействиям с трекером, его можно быстро получить с помощью метода Accelerator.get_tracker(). Достаточно передать строку, соответствующую атрибуту .name трекера, и он вернёт этот трекер в основном процессе.

В этом примере показано, как это сделать с помощью wandb:

wandb_tracker = accelerator.get_tracker("wandb")

Оттуда вы можете взаимодействовать с объектом run в wandb как обычно:

wandb_tracker.log_artifact(some_artifact_to_log)

Trackers built in Accelerate will automatically execute on the correct process, so if a tracker is only meant to be ran on the main process it will do so automatically.

Если вы хотите полностью удалить обёртку Accelerate, вы можете достичь того же результата с помощью:

wandb_tracker = accelerator.get_tracker("wandb", unwrap=True)
if accelerator.is_main_process:
    wandb_tracker.log_artifact(some_artifact_to_log)

Когда обёртка не может работать

Если у библиотеки есть API, не следующее строгому формату .log с общим словарём, как в Neptune.AI, логирование можно выполнять вручную в рамках оператора if accelerator.is_main_process:

  from accelerate import Accelerator
+ import neptune

  accelerator = Accelerator()
+ run = neptune.init_run(...)

  my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
  device = accelerator.device
  my_model.to(device)

  for iteration in config["num_iterations"]:
      for batch in my_training_dataloader:
          my_optimizer.zero_grad()
          inputs, targets = batch
          inputs = inputs.to(device)
          targets = targets.to(device)
          outputs = my_model(inputs)
          loss = my_loss_function(outputs, targets)
          total_loss += loss
          accelerator.backward(loss)
          my_optimizer.step()
+         if accelerator.is_main_process:
+             run["logs/training/batch/loss"].log(loss)