实验跟踪器

市面上有大量实验跟踪 API 可用，但要在多进程环境中让它们全部正常工作往往十分复杂。 Accelerate 提供了一个通用的跟踪 API，您可以通过 Accelerator.log() 在脚本运行期间记录有用的信息。

集成追踪器

目前 Accelerate 开箱即用地支持八种跟踪器：

TensorBoard
WandB
Trackio
CometML
目标
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()，以便所有 tracker 能够执行其结束功能（如果有的话）。

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。当 ProjectConfiguration 数据类中还有其他需要组合的配置时，project_dir 参数非常有用。例如，您可以将 TensorBoard 数据保存到 project_dir，而将其他所有内容记录在 [~utils.ProjectConfiguration 的 logging_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)

实现自定义跟踪器

要在 Accelerator 中实现一个新的跟踪器，可以通过实现 GeneralTracker 类来创建。每个跟踪器都必须实现三个函数并拥有三个属性：

__init__:
- 应存储 run_name 并初始化集成库的 tracker API。
- 如果跟踪器将数据存储在本地（例如 TensorBoard），则可以添加 logging_dir 参数。
store_init_configuration:
- 应接收一个values字典，并将其存储为一次性实验配置
log:
- 应接收一个values字典和一个step，并将它们记录到运行中
name (str):
- 跟踪器的唯一字符串名称，例如用于 wandb 跟踪器的 "wandb"。
- 这将专门用于与此跟踪器进行交互
requires_logging_directory (bool):
- 此特定跟踪器是否需要logging_dir，以及它是否使用了该目录。
tracker:
- 这应实现为一个 @property 函数
- 应返回该库使用的内部跟踪机制，例如 wandb 的 run 对象。

每个方法还应使用 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")

从那里开始，您可以像往常一样与 wandb 的 run 对象进行交互：

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)

Update on GitHub