メインコンテンツへスキップ
Try in Colab Hugging Face Transformers ライブラリを使用すると、BERT のような最先端の NLP モデルや、混合精度トレーニング、勾配チェックポイントなどのトレーニング手法を簡単に利用できます。W&B インテグレーション を活用することで、その使いやすさを損なうことなく、リッチで柔軟な 実験管理 と モデルのバージョン管理を、インタラクティブで一元化されたダッシュボードに追加できます。

数行のコードで高度なロギングを実現

os.environ["WANDB_PROJECT"] = "<my-amazing-project>"  # W&B プロジェクト名を指定
os.environ["WANDB_LOG_MODEL"] = "checkpoint"  # すべてのモデルチェックポイントをログに記録

from transformers import TrainingArguments, Trainer

args = TrainingArguments(..., report_to="wandb")  # W&B ロギングを有効化
trainer = Trainer(..., args=args)
HuggingFace dashboard
すぐに実行可能なコードを確認したい場合は、こちらの Google Colab をご覧ください。

はじめに:実験を追跡する

サインアップと API キーの作成

APIキー は、お使いのマシンを W&B に対して認証するために使用されます。ユーザープロファイルから APIキー を生成できます。
For a more streamlined approach, create an API key by going directly to User Settings. Copy the newly created API key immediately and save it in a secure location such as a password manager.
  1. 右上隅にあるユーザープロファイルアイコンをクリックします。
  2. User Settings を選択し、API Keys セクションまでスクロールします。

wandb ライブラリのインストールとログイン

ローカルに wandb ライブラリをインストールしてログインするには:
  1. WANDB_API_KEY 環境変数 を APIキー に設定します。
    export WANDB_API_KEY=<your_api_key>
    
  2. wandb ライブラリをインストールしてログインします。
    pip install wandb
    
    wandb login
    
W&B を初めて使用する場合は、クイックスタート も併せてご確認ください。

プロジェクトに名前を付ける

W&B の Projects は、関連する Runs からログ記録されたすべてのチャート、データ、モデルが保存される場所です。プロジェクトに名前を付けることで、作業を整理し、1つのプロジェクトに関するすべての情報を1か所にまとめて管理できます。 Run をプロジェクトに追加するには、環境変数 WANDB_PROJECT にプロジェクト名を設定するだけです。WandbCallback がこのプロジェクト名の環境変数を読み取り、Run のセットアップ時に使用します。
WANDB_PROJECT=amazon_sentiment_analysis
プロジェクト名は必ず Trainer を初期化する に設定してください。
プロジェクト名が指定されていない場合、デフォルトのプロジェクト名は huggingface になります。

トレーニング Run を W&B にログ記録する

Trainer のトレーニング引数を定義する際、コード内またはコマンドラインから 最も重要なステップ は、W&B でのロギングを有効にするために report_to"wandb" に設定することです。 TrainingArguments 内の logging_steps 引数により、トレーニング中にメトリクスを W&B にプッシュする頻度を制御できます。また、run_name 引数を使用して、W&B 上のトレーニング Run に名前を付けることも可能です。 設定はこれだけです。これで、トレーニング中の損失、評価メトリクス、モデルのトポロジー、勾配が W&B にログ記録されるようになります。
python run_glue.py \     # Python スクリプトを実行
  --report_to wandb \    # W&B へのロギングを有効化
  --run_name bert-base-high-lr \   # W&B Run の名前(任意)
  # その他のコマンドライン引数
TensorFlow をお使いですか? PyTorch の Trainer を TensorFlow 用の TFTrainer に置き換えるだけで同様に動作します。

モデルのチェックポイント保存を有効にする

Artifacts を使用すると、最大 100GB までのモデルやデータセットを無料で保存でき、W&B の Registry を利用できるようになります。Registry を使えば、モデルの登録、探索、評価、ステージングへの準備、プロダクション環境へのデプロイが可能になります。 Hugging Face のモデルチェックポイントを Artifacts にログ記録するには、環境変数 WANDB_LOG_MODEL を以下のいずれかに設定します:
  • checkpoint: TrainingArgumentsargs.save_steps ごとにチェックポイントをアップロードします。
  • end: load_best_model_at_end も設定されている場合、トレーニング終了時にモデルをアップロードします。
  • false: モデルをアップロードしません。
WANDB_LOG_MODEL="checkpoint"
これ以降に初期化されるすべての Transformers Trainer は、モデルを W&B プロジェクトにアップロードします。ログ記録されたモデルチェックポイントは Artifacts UI から確認でき、完全な モデルリネージ も含まれます(UI でのモデルチェックポイントの例は こちら で確認できます)。
デフォルトでは、WANDB_LOG_MODELend に設定されている場合は model-{run_id}checkpoint に設定されている場合は checkpoint-{run_id} という名前で W&B Artifacts に保存されます。 ただし、TrainingArgumentsrun_name を渡した場合、モデルは model-{run_name} または checkpoint-{run_name} として保存されます。

W&B Registry

チェックポイントを Artifacts にログ記録した後は、最良のモデルチェックポイントを登録し、Registry を通じてチーム全体で一元管理できます。Registry を使用すると、タスクごとに最適なモデルを整理し、モデルのライフサイクルを管理し、ML ライフサイクル全体を追跡・監査し、ダウンストリームのアクションを オートメーション 化できます。 モデル Artifact のリンク方法については、Registry を参照してください。

トレーニング中の評価出力の可視化

トレーニング中や評価中のモデル出力を可視化することは、モデルがどのように学習しているかを真に理解するために不可欠です。 Transformers Trainer のコールバックシステムを使用すると、モデルのテキスト生成出力やその他の予測結果などの追加データを W&B Tables にログ記録できます。 トレーニング中の評価出力を W&B テーブルにログ記録する方法の完全なガイドについては、以下の カスタムロギングセクション を参照してください。
評価出力を表示する W&B テーブル

W&B Run の終了(ノートブックのみ)

トレーニングが Python スクリプトにカプセル化されている場合、W&B Run はスクリプトの終了とともに完了します。 Jupyter や Google Colab ノートブックを使用している場合は、run.finish() を呼び出してトレーニングの終了を明示的に伝える必要があります。
run = wandb.init()
trainer.train()  # トレーニングを開始し W&B にログを記録

# トレーニング後の分析、テスト、その他のログ記録コード

run.finish()

結果の可視化

トレーニング結果をログ記録した後は、W&B Dashboard で動的に結果を探索できます。数十の Run を一度に比較したり、興味深い発見をズームアップしたり、柔軟でインタラクティブな可視化機能を使って複雑なデータから洞察を引き出したりすることが簡単にできます。

高度な機能と FAQ

最良のモデルを保存するには?

load_best_model_at_end=True を含む TrainingArgumentsTrainer に渡すと、W&B は最もパフォーマンスの高いモデルチェックポイントを Artifacts に保存します。 モデルチェックポイントを Artifacts として保存すると、それらを Registry にプロモートできます。Registry では以下が可能です:
  • ML タスクごとに最適なモデルバージョンを整理。
  • モデルを一元管理し、チームで共有。
  • プロダクション用のモデルをステージング、またはさらなる評価のためにブックマーク。
  • ダウンストリームの CI/CD プロセスをトリガー。

保存したモデルをロードするには?

WANDB_LOG_MODEL を使用してモデルを W&B Artifacts に保存した場合、追加のトレーニングや推論実行のためにモデルの重みをダウンロードできます。以前使用したものと同じ Hugging Face アーキテクチャーにそれらをロードし直すだけです。
# 新しい run を作成
with wandb.init(project="amazon_sentiment_analysis") as run:
    # アーティファクトの名前とバージョンを指定
    my_model_name = "model-bert-base-high-lr:latest"
    my_model_artifact = run.use_artifact(my_model_name)

    # モデルの重みをフォルダーにダウンロードし、パスを返す
    model_dir = my_model_artifact.download()

    # そのフォルダーから同じモデルクラスを使用して 
    # Hugging Face モデルをロード
    model = AutoModelForSequenceClassification.from_pretrained(
        model_dir, num_labels=num_labels
    )

    # 追加のトレーニングや推論の実行

チェックポイントからトレーニングを再開するには?

WANDB_LOG_MODEL='checkpoint' を設定していた場合、model_dirTrainingArgumentsmodel_name_or_path 引数として使用し、Trainerresume_from_checkpoint=True を渡すことでトレーニングを再開できます。
last_run_id = "xxxxxxxx"  # wandb ワークスペースから run_id を取得

# run_id を指定して wandb run を再開
with wandb.init(
    project=os.environ["WANDB_PROJECT"],
    id=last_run_id,
    resume="must",
) as run:
    # Run に Artifact を接続
    my_checkpoint_name = f"checkpoint-{last_run_id}:latest"
    my_checkpoint_artifact = run.use_artifact(my_model_name)

    # チェックポイントをフォルダーにダウンロードし、パスを返す
    checkpoint_dir = my_checkpoint_artifact.download()

    # モデルとトレーナーを再初期化
    model = AutoModelForSequenceClassification.from_pretrained(
        "<model_name>", num_labels=num_labels
    )
    # トレーニング引数の設定
    training_args = TrainingArguments()

    trainer = Trainer(model=model, args=training_args)

    # チェックポイントディレクトリを使用してトレーニングを再開
    trainer.train(resume_from_checkpoint=checkpoint_dir)

トレーニング中に評価サンプルをログ記録・表示するには

Transformers Trainer を介した W&B へのロギングは、Transformers ライブラリ内の WandbCallback によって処理されます。Hugging Face のロギングをカスタマイズする必要がある場合は、WandbCallback をサブクラス化し、Trainer クラスのメソッドを活用する機能を追加することで、このコールバックを変更できます。 以下は、この新しいコールバックを HF Trainer に追加する一般的なパターンです。さらにその下には、評価出力を W&B テーブルにログ記録するための完全なコード例があります。
# 通常通り Trainer をインスタンス化
trainer = Trainer()

# 新しいロギングコールバックをインスタンス化し、Trainer オブジェクトを渡す
evals_callback = WandbEvalsCallback(trainer, tokenizer, ...)

# コールバックを Trainer に追加
trainer.add_callback(evals_callback)

# 通常通りトレーニングを開始
trainer.train()

トレーニング中に評価サンプルを表示する

以下のセクションでは、WandbCallback をカスタマイズして、トレーニング中にモデルの予測を実行し、評価サンプルを W&B テーブルにログ記録する方法を示します。Trainer コールバックの on_evaluate メソッドを使用して、eval_steps ごとに実行します。 ここでは、トークナイザーを使用してモデル出力から予測とラベルをデコードする decode_predictions 関数を作成しました。 次に、予測とラベルから pandas DataFrame を作成し、DataFrame に epoch 列を追加します。 最後に、DataFrame から wandb.Table を作成し、それを W&B にログ記録します。 さらに、freq エポックごとに予測をログ記録することで、ログの頻度を制御できます。 注意: 通常の WandbCallback とは異なり、このカスタムコールバックは Trainer の初期化中ではなく、Trainer がインスタンス化された に追加する必要があります。 これは、初期化時に Trainer インスタンスがコールバックに渡されるためです。
from transformers.integrations import WandbCallback
import pandas as pd


def decode_predictions(tokenizer, predictions):
    labels = tokenizer.batch_decode(predictions.label_ids)
    logits = predictions.predictions.argmax(axis=-1)
    prediction_text = tokenizer.batch_decode(logits)
    return {"labels": labels, "predictions": prediction_text}


class WandbPredictionProgressCallback(WandbCallback):
    """トレーニング中にモデルの予測をログ記録するカスタム WandbCallback。

    このコールバックは、トレーニング中の各ロギングステップでモデルの予測とラベルを
    wandb.Table にログ記録します。これにより、トレーニングの進行に合わせて
    モデルの予測を可視化できます。

    Attributes:
        trainer (Trainer): Hugging Face Trainer インスタンス。
        tokenizer (AutoTokenizer): モデルに関連付けられたトークナイザー。
        sample_dataset (Dataset): 予測生成用の検証データセットのサブセット。
        num_samples (int, optional): 予測生成用に検証データセットから選択するサンプル数。デフォルトは 100。
        freq (int, optional): ロギングの頻度。デフォルトは 2。
    """

    def __init__(self, trainer, tokenizer, val_dataset, num_samples=100, freq=2):
        super().__init__()
        self.trainer = trainer
        self.tokenizer = tokenizer
        self.sample_dataset = val_dataset.select(range(num_samples))
        self.freq = freq

    def on_evaluate(self, args, state, control, **kwargs):
        super().on_evaluate(args, state, control, **kwargs)
        # `freq` エポックごとに予測をログ記録することで頻度を制御
        if state.epoch % self.freq == 0:
            # 予測を生成
            predictions = self.trainer.predict(self.sample_dataset)
            # 予測とラベルをデコード
            predictions = decode_predictions(self.tokenizer, predictions)
            # 予測を wandb.Table に追加
            predictions_df = pd.DataFrame(predictions)
            predictions_df["epoch"] = state.epoch
            records_table = self._wandb.Table(dataframe=predictions_df)
            # テーブルを wandb にログ記録
            self._wandb.log({"sample_predictions": records_table})


# まず、Trainer をインスタンス化
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_datasets["train"],
    eval_dataset=lm_datasets["validation"],
)

# WandbPredictionProgressCallback をインスタンス化
progress_callback = WandbPredictionProgressCallback(
    trainer=trainer,
    tokenizer=tokenizer,
    val_dataset=lm_dataset["validation"],
    num_samples=10,
    freq=2,
)

# コールバックを trainer に追加
trainer.add_callback(progress_callback)
より詳細な例については、こちらの Colab を参照してください。

利用可能な追加の W&B 設定は?

環境変数を設定することで、Trainer でログ記録される内容をさらに詳細に設定できます。W&B 環境変数の全リストは こちらで見ることができます
環境変数用途
WANDB_PROJECTプロジェクトに名前を付けます(デフォルトは huggingface
WANDB_LOG_MODEL

モデルチェックポイントを W&B Artifact としてログ記録します(デフォルトは false

  • false (デフォルト): チェックポイントを記録しません
  • checkpoint: args.save_steps ごとにチェックポイントがアップロードされます。
  • end: トレーニング終了時に最終的なモデルチェックポイントがアップロードされます。
WANDB_WATCH

モデルの勾配、パラメータ、またはその両方をログ記録するか設定します

  • false (デフォルト): 勾配やパラメータのログを記録しません
  • gradients: 勾配のヒストグラムをログ記録します
  • all: 勾配とパラメータの両方のヒストグラムをログ記録します
WANDB_DISABLEDロギングを完全に無効にするには true に設定します(デフォルトは false
WANDB_QUIET.標準出力に記録されるステートメントを重要なものだけに制限するには true に設定します(デフォルトは false
WANDB_SILENTwandb によって出力されるテキストを非表示にするには true に設定します(デフォルトは false
WANDB_WATCH=all
WANDB_SILENT=true

wandb.init をカスタマイズするには?

Trainer が使用する WandbCallback は、Trainer の初期化時にバックグラウンドで wandb.init を呼び出します。あるいは、Trainer を初期化する前に wandb.init を呼び出して手動で Run を設定することもできます。これにより、W&B Run の設定を完全に制御できるようになります。 init に渡す設定の例を以下に示します。wandb.init() の詳細については、wandb.init() リファレンス を参照してください。
wandb.init(
    project="amazon_sentiment_analysis",
    name="bert-base-high-lr",
    tags=["baseline", "high-lr"],
    group="bert",
)

その他のリソース

Transformers と W&B に関連する 6 つの記事をご紹介します。

ヘルプの取得や機能リクエスト

Hugging Face W&B インテグレーションに関する問題、質問、機能リクエストについては、Hugging Face フォーラムのこのスレッド に投稿するか、Hugging Face Transformers GitHub リポジトリ で Issue を作成してください。