メインコンテンツへスキップ
Hugging Face Accelerate は、同じ PyTorch コードをあらゆる分散設定で実行できるようにし、大規模なモデル学習と推論を簡素化するためのライブラリです。 Accelerate には W&B Tracker が含まれており、以下でその使い方を説明します。Hugging Face における Accelerate Trackers の詳細も参照できます。

Accelerate でロギングを開始する

Accelerate と W&B の連携を始めるには、次の疑似コードを参考にしてください。
from accelerate import Accelerator

# Acceleratorオブジェクトにwandbでログするよう指示する
accelerator = Accelerator(log_with="wandb")

# wandb runを初期化し、wandbパラメータと設定情報を渡す
accelerator.init_trackers(
    project_name="my_project", 
    config={"dropout": 0.1, "learning_rate": 1e-2}
    init_kwargs={"wandb": {"entity": "my-wandb-team"}}
    )

...

# `accelerator.log`を呼び出してwandbにログする。`step`は省略可能
accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=global_step)


# wandbトラッカーが正しく終了することを確認する
accelerator.end_training()
もう少し詳しく説明すると、次のことを行う必要があります。
  1. Accelerator クラスを初期化するときに log_with="wandb" を渡す
  2. init_trackers メソッドを呼び出し、次の引数を渡す:
  • project_name にプロジェクト名
  • init_kwargs にネストされた dict を渡すことで、wandb.init() に渡したい任意のパラメータ
  • config に、wandb の run にログしたいその他の実験設定情報
  1. .log メソッドを使って Weights & Biases にログを送る(step 引数は省略可能)
  2. 学習が完了したら .end_training を呼び出す

W&B トラッカーにアクセスする

W&B トラッカーにアクセスするには、Accelerator.get_tracker() メソッドを使用します。トラッカーの .name 属性に対応する文字列を渡すと、メインプロセスのトラッカーが返されます。
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)
with accelerator.on_main_process:
    wandb_tracker.log_artifact(some_artifact_to_log)

Accelerate 記事

以下は、Accelerate に関するおすすめの記事です。