加速文档

实验追踪器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验

开始使用

实验追踪器

有很多可用的实验追踪 API,但是让它们在多进程环境中协同工作通常很复杂。Accelerate 提供了一个通用追踪 API,可用于通过 Accelerator.log() 在您的脚本运行期间记录有用的项目。

集成追踪器

目前,Accelerate 支持开箱即用的七个追踪器。

  • TensorBoard
  • WandB
  • CometML
  • Aim
  • MLFlow
  • ClearML
  • DVCLive

要使用其中任何一个,请在 Accelerate 中将选定的类型传递给 log_with 参数。

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 = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
device = accelerator.device
my_model.to(device)

for iteration in config["num_iterations"]:
    for step, 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)
        accelerator.backward(loss)
        my_optimizer.step()
        accelerator.log({"training_loss": loss}, step=step)
accelerator.end_training()

如果追踪器需要一个目录来保存数据,例如 TensorBoard,则将目录路径传递给 project_dir。当需要与 ProjectConfiguration 数据类中的其他配置组合时,project_dir 参数非常有用。例如,您可以将 TensorBoard 数据保存到 project_dir,而所有其他内容都可以记录在 [~utils.ProjectConfigurationlogging_dir 参数中。

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)

实现自定义追踪器

要实现一个在 Accelerate 中使用的新的追踪器,可以通过实现 GeneralTracker 类来创建一个新的追踪器。每个追踪器都必须实现三个函数,并具有三个属性。

  • __init__:

    • 应该存储一个 run_name 并初始化集成库的追踪器 API。
    • 如果追踪器在本地存储其数据(例如 TensorBoard),则可以添加一个 logging_dir 参数。
  • store_init_configuration:

    • 应该接收一个 values 字典并将它们存储为一次性实验配置。
  • log:

    • 应该接收一个 values 字典和一个 step,并将它们记录到运行中。
  • name (str)

    • 追踪器的唯一字符串名称,例如 "wandb" 代表 wandb 追踪器。
    • 这将用于与该特定追踪器进行交互。
  • requires_logging_directory (bool)

    • 此特定追踪器是否需要 logging_dir 以及它是否使用 logging_dir
  • tracker:

    • 这应该实现为一个 @property 函数。
    • 应该返回库使用的内部追踪机制,例如 wandbrun 对象。

如果记录器仅应在主进程上执行,则每个方法还应利用 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"])

访问内部追踪器

如果可能需要直接与追踪器进行一些自定义交互,您可以使用 Accelerator.get_tracker() 方法快速访问一个追踪器。只需传入与追踪器的 .name 属性相对应的字符串,它将在主进程上返回该追踪器。

此示例展示了如何使用 wandb 进行操作。

wandb_tracker = accelerator.get_tracker("wandb")

从那里您可以像平常一样与 wandbrun 对象进行交互。

wandb_run.log_artifact(some_artifact_to_log)
在 Accelerate 中构建的追踪器将自动在正确的进程上执行,因此,如果追踪器只打算在主进程上运行,它将自动执行。

如果您想完全删除 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.new as neptune

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

  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)
< > 在 GitHub 上更新