メインコンテンツへスキップ
Hugging Face Accelerate は、同じ PyTorch コードをあらゆる分散構成で実行できるようにし、大規模なモデルのトレーニングと推論を簡素化するライブラリです。 Accelerate には W&B トラッカーが含まれており、以下でその使い方を紹介します。また、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 method を呼び出し、次を渡します。
  • project_name でプロジェクト名を渡す
  • wandb.init() に渡したい任意のパラメーターを、ネストした dict として init_kwargs に渡す
  • wandb run にログしたいその他の実験設定情報を config で渡す
  1. Weights & Biases にログするには wandb.Run.log() method を使用します。step 引数は省略可能です
  2. トレーニングが完了したら .end_training() を呼び出します

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

W&Bトラッカーにアクセスするには、Accelerator.get_tracker() method を使用します。トラッカーの .name 属性に対応する文字列を渡すと、main プロセス上のトラッカーが返されます。 
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 の記事もぜひご覧ください