實驗跟蹤器

有大量的實驗跟蹤 API 可用，但讓它們在多處理環境中協同工作通常很複雜。Accelerate 提供了一個通用的跟蹤 API，可以透過 Accelerator.log() 在指令碼中記錄有用的專案。

整合跟蹤器

目前，Accelerate 開箱即用地支援八種跟蹤器：

TensorBoard
WandB
Trackio
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 = 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 並初始化整合庫的跟蹤器 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"])

訪問內部跟蹤器

如果希望直接與某個跟蹤器進行一些自定義互動，可以使用 Accelerator.get_tracker() 方法快速訪問它。只需傳入與跟蹤器 .name 屬性對應的字串，它就會在主程序上返回該跟蹤器。

這個例子展示瞭如何用 wandb 來做到這一點。

wandb_tracker = accelerator.get_tracker("wandb")

從那裡你可以像往常一樣與 wandb 的 run 物件進行互動。

wandb_tracker.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

  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)

< > 在 GitHub 上更新