LlamaFactory
算力平台:

实验跟踪器 ​

有许多实验跟踪 API 可用，但在多进程环境中让它们全部正常工作通常会很复杂。 Accelerate 提供了一个通用的跟踪 API，可以通过 [Accelerator.log] 在脚本中记录有用的信息。

集成的跟踪器 ​

目前 Accelerate 自带支持七种跟踪器：

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

要使用其中任何一种，只需将选定的类型传递给 [Accelerate] 中的 log_with 参数：

python
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] 来设置你的项目，并可能添加任何实验的超参数以进行记录：

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

当你准备记录任何数据时，应使用 [Accelerator.log]。 还可以传递一个 step 参数，以便将数据与训练循环中的特定步骤关联起来。

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

完成训练后，务必运行 [Accelerator.end_training]，以便所有跟踪器能够执行它们的结束功能（如果有的话）。

python
accelerator.end_training()

一个完整的示例如下：

python
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。project_dir 参数在需要与其他配置结合使用时非常有用，例如在 [~utils.ProjectConfiguration] 数据类中。例如，你可以将 TensorBoard 数据保存到 project_dir，而其他所有内容可以记录在 logging_dir 参数中。

python
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，并将其记录到运行中。

  • name (str):

    • 跟踪器的唯一字符串名称，例如 wandb 跟踪器的 "wandb"。
    • 这将用于专门与该跟踪器进行交互。

  • requires_logging_directory (bool):

    • 是否需要 logging_dir 以及该特定跟踪器是否使用它。

  • tracker:

    • 应实现为一个 @property 函数。
    • 应返回库使用的内部跟踪机制，例如 wandb 的 run 对象。

每个方法还应使用 [state.PartialState] 类，例如，如果日志记录器仅应在主进程中执行。

以下是一个简短的示例，展示了与 Weights and Biases 的集成，仅包含相关的信息并在主进程中进行日志记录：

python
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 一起使用：

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

这些也可以与现有的追踪器混合使用，包括与 "all" 一起使用：

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

访问内部跟踪器 ​

如果需要直接与跟踪器进行一些自定义交互，可以使用 [Accelerator.get_tracker] 方法快速访问跟踪器。只需传入与跟踪器的 .name 属性对应的字符串，它将在主进程中返回该跟踪器。

以下示例展示了如何使用 wandb 进行操作：

python
wandb_tracker = accelerator.get_tracker("wandb")

从那里你可以像平常一样与 wandb 的 run 对象进行交互：

python
wandb_run.log_artifact(some_artifact_to_log)

如果你想要完全移除 Accelerate 的包装，你可以通过以下方式实现相同的效果：

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

当包装器无法工作时 ​

如果一个库的 API 没有遵循像 Neptune.AI 那样的严格 .log 方法，可以在 if accelerator.is_main_process 语句下手动进行日志记录：

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

Copyright © 2024-2025 llamafactory.cn