Seguimiento de experimentos

Existe un gran número de APIs de seguimiento de experimentos disponibles; sin embargo, lograr que todas funcionen en un entorno de multiprocesamiento puede ser complejo con frecuencia. Accelerate proporciona una API de seguimiento general que se puede utilizar para registrar elementos útiles durante la ejecución de su script a través de Accelerator.log()

Seguimientos integrados

Actualmente, Accelerate admite ocho trackers de forma nativa:

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

Para usar cualquiera de ellos, pase el tipo o los tipos seleccionados al parámetro log_with en 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])

Al inicio de su experimento, debe usar Accelerator.init_trackers() para configurar su proyecto y, potencialmente, agregar los hiperparámetros del experimento que se registrarán:

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

Cuando esté listo para registrar cualquier dato, debe usarse Accelerator.log(). También se puede pasar un step para correlacionar los datos con un paso específico del bucle de entrenamiento.

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

Una vez que hayas terminado el entrenamiento, asegúrate de ejecutar Accelerator.end_training() para que todos los trackers puedan ejecutar sus funcionalidades de finalización, si las tienen.

accelerator.end_training()

A continuación se muestra un ejemplo completo:

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

Si un tracker requiere un directorio para guardar los datos, como TensorBoard, entonces pase la ruta del directorio a project_dir. El parámetro project_dir es útil cuando hay otras configuraciones que deben combinarse en la clase de datos ProjectConfiguration. Por ejemplo, puede guardar los datos de TensorBoard en project_dir y registrar todo lo demás en el 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)

Implementación de rastreadores personalizados

Para implementar un nuevo tracker para usar en Accelerator, se puede crear uno nuevo implementando la clase GeneralTracker. Cada tracker debe implementar tres funciones y tener tres propiedades:

• __init__:

  • Debe almacenar un run_name e inicializar la API del tracker de la biblioteca integrada.
  • Si un tracker almacena sus datos localmente (como TensorBoard), se puede añadir un parámetro logging_dir.

• store_init_configuration:

  • Debe recibir un diccionario values y almacenarlos como una configuración de experimento de un solo uso

• log:

  • Debe recibir un diccionario values y un step, y debe registrarlos en la ejecución

• name (str):

  • Un nombre de cadena único para el tracker, como "wandb" para el tracker de wandb.
  • Esto se utilizará para interactuar específicamente con este tracker

• requires_logging_directory (bool):

  • Si se necesita un logging_dir para este tracker en particular y si lo utiliza.

• tracker:

  • Esto debería implementarse como una función @property
  • Debe devolver el mecanismo de seguimiento interno que utiliza la biblioteca, como el objeto run para wandb.

Cada método también debe utilizar la clase state.PartialState si, por ejemplo, el registrador solo debe ejecutarse en el proceso principal.

A continuación se muestra un ejemplo breve con una integración con Weights and Biases, que contiene solo la información relevante y registra únicamente en el proceso 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)

Cuando esté listo para crear su objeto Accelerator, pase una instancia de su tracker a Accelerator.log_with para que se utilice automáticamente con la API:

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

También se pueden combinar con los rastreadores existentes, incluido "all":

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

Registrar un rastreador personalizado por nombre

En lugar de pasar una instancia, puede registrar una clase de tracker personalizada para que pueda seleccionarse por su name, al igual que los trackers integrados. Esto resulta útil cuando los trackers a utilizar se leen de un archivo de configuración o de la línea de comandos:

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

Accediendo al rastreador interno

Si se desean interacciones personalizadas directamente con un tracker, se puede acceder rápidamente a uno mediante el método Accelerator.get_tracker(). Basta con pasar la cadena de texto correspondiente al atributo .name de un tracker y este se devolverá en el proceso principal.

Este ejemplo muestra cómo hacerlo con wandb:

wandb_tracker = accelerator.get_tracker("wandb")

A partir de ahí, puedes interactuar con el objeto run de wandb de forma habitual:

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.

Si deseas eliminar por completo el envoltorio de Accelerate, puedes lograr el mismo resultado con:

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

Cuando un envoltorio no puede funcionar

Si una biblioteca tiene una API que no sigue un estricto .log con un diccionario general, como Neptune.AI, el registro se puede realizar manualmente bajo una declaración 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)