Rastreadores de experimentos

Há um grande número de APIs de rastreamento de experimentos disponíveis, no entanto, fazer com que todas funcionem em um ambiente de multiprocessamento pode ser frequentemente complexo. O Accelerate fornece uma API de rastreamento geral que pode ser usada para registrar itens úteis durante a execução do seu script por meio de Accelerator.log()

Rastreadores Integrados

Atualmente, o Accelerate suporta oito rastreadores nativamente:

• TensorBoard
• WandB
• Trackio
• CometML
• Objetivo
• MLFlow
• ClearML
• DVCLive

Para usar qualquer um deles, passe o(s) tipo(s) selecionado(s) para o parâmetro log_with no 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])

No início do seu experimento, Accelerator.init_trackers() deve ser usado para configurar seu projeto e, potencialmente, adicionar quaisquer hiperparâmetros do experimento a serem registrados:

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

Quando estiver pronto para registrar qualquer dado, Accelerator.log() deve ser usado. Um step também pode ser passado para correlacionar os dados com uma etapa específica no loop de treinamento.

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

Depois de terminar o treinamento, certifique-se de executar Accelerator.end_training() para que todos os trackers possam executar suas funcionalidades de finalização, caso as tenham.

accelerator.end_training()

Um exemplo completo está abaixo:

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()

Se um tracker exigir um diretório para salvar os dados, como o TensorBoard, passe o caminho do diretório para project_dir. O parâmetro project_dir é útil quando há outras configurações a serem combinadas na classe de dados ProjectConfiguration. Por exemplo, você pode salvar os dados do TensorBoard em project_dir e registrar tudo o mais no parâmetro logging_dir de [~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)

Implementando rastreadores personalizados

Para implementar um novo tracker para ser usado em Accelerator, um novo pode ser criado por meio da implementação da classe GeneralTracker. Cada tracker deve implementar três funções e ter três propriedades:

• __init__:

  • Deve armazenar um run_name e inicializar a API do tracker da biblioteca integrada.
  • Se um tracker armazena seus dados localmente (como o TensorBoard), um parâmetro logging_dir pode ser adicionado.

• store_init_configuration:

  • Deve receber um dicionário values e armazená-lo como uma configuração de experimento de uso único

• log:

  • Deve receber um dicionário values e um step, e deve registrá-los na execução

• name (str):

  • Um nome de string único para o tracker, como "wandb" para o tracker wandb.
  • Isso será usado para interagir especificamente com esse rastreador

• requires_logging_directory (bool):

  • Se um logging_dir é necessário para este tracker específico e se ele utiliza um.

• tracker:

  • Isso deve ser implementado como uma função @property
  • Deve retornar o mecanismo de rastreamento interno usado pela biblioteca, como o objeto run para o wandb.

Cada método também deve utilizar a classe state.PartialState se o logger, por exemplo, deve ser executado apenas no processo principal.

Um breve exemplo pode ser visto abaixo com uma integração com Weights and Biases, contendo apenas as informações relevantes e registrando apenas no processo principal:

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)

Quando estiver pronto para criar seu objeto Accelerator, passe uma instância do seu tracker para Accelerator.log_with para que ele seja usado automaticamente com a API:

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

Estes também podem ser misturados com rastreadores existentes, incluindo com "all":

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

Registrando um rastreador personalizado por nome

Em vez de passar uma instância, você pode registrar uma classe de rastreador personalizada para que ela possa ser selecionada pelo seu name, assim como os rastreadores integrados. Isso é útil quando os rastreadores a serem usados são lidos de um arquivo de configuração ou da linha de comando:

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")

Acessando o rastreador interno

Se algumas interações personalizadas com um tracker forem desejadas diretamente, você pode acessá-lo rapidamente usando o método Accelerator.get_tracker(). Basta passar a string correspondente ao atributo .name de um tracker e ele retornará esse tracker no processo principal.

Este exemplo mostra como fazer isso com wandb:

wandb_tracker = accelerator.get_tracker("wandb")

A partir daí, você pode interagir com o objeto run do wandb normalmente:

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.

Se você deseja remover completamente o encapsulamento do Accelerate, pode alcançar o mesmo resultado com:

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

Quando um wrapper não pode funcionar

Se uma biblioteca possui uma API que não segue um rigoroso .log com um dicionário geral, como o Neptune.AI, o registro pode ser feito manualmente sob uma declaração 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)