メインコンテンツへスキップ
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 ダッシュボード
そのまま実行できるコードから始めたい場合は、この Google Colab を参照してください。

Experiments のトラッキングを始める

サインアップしてAPIキーを作成する

APIキーは、あなたのマシンをW&Bに対して認証するためのものです。APIキーはユーザープロフィールから作成できます。
よりシンプルに行うには、User Settings に直接アクセスしてAPIキーを作成してください。作成したばかりのAPIキーをすぐにコピーし、パスワードマネージャーなどの安全な場所に保管してください。
  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 のプロジェクトは、関連する run から記録されたすべてのチャート、データ、モデルが保存される場所です。プロジェクトに名前を付けることで、作業を整理し、1 つのプロジェクトに関するすべての情報を 1 か所にまとめておけます。 run をプロジェクトに追加するには、WANDB_PROJECT 環境変数にプロジェクト名を設定します。WandbCallback はこのプロジェクト名の環境変数を自動的に読み取り、run をセットアップするときに使用します。
WANDB_PROJECT=amazon_sentiment_analysis
Trainer を初期化する 前に プロジェクト名を設定しておいてください。
プロジェクト名を指定しない場合、プロジェクト名はデフォルトで huggingface になります。

学習 run を W&B に記録する

コード内またはコマンドラインから Trainer の学習引数を指定する際の 最も重要な手順 は、W&B へのロギングを有効にするために report_to"wandb" に設定することです。 TrainingArgumentslogging_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 を次の いずれか 1 つ に設定します:
  • 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 のライフサイクル全体をトラッキングおよび監査し、後続のアクションを自動化できます。 モデルのアーティファクトをリンクする方法については、Registry を参照してください。

学習中に評価出力を可視化する

学習や評価の最中にモデルの出力を可視化することは、モデルがどのように学習されているかを深く理解するうえで、多くの場合不可欠です。 Transformers の Trainer のコールバックシステムを使用すると、モデルのテキスト生成出力やその他の予測などの、役立つ追加データを W&B にログとして記録し、W&B Tables に保存できます。 以下のカスタムロギングのセクションでは、学習中に評価出力をログし、このような W&B Table に記録する方法についての詳細なガイドを確認できます。
評価出力を含む W&B Table を表示しているスクリーンショット

W&B の run を終了する(ノートブックのみ)

学習が Python スクリプトとしてまとめられている場合は、スクリプトが終了すると同時に W&B の run も終了します。 Jupyter Notebook または Google Colab を使用している場合は、run.finish() を呼び出して、学習が完了したことを W&B に明示的に伝える必要があります。
run = wandb.init()
trainer.train()  # 学習を開始してW&Bにログを記録する

# 学習後の分析、テスト、その他のログ記録コード

run.finish()

結果を可視化する

学習結果をログしたら、W&B Dashboard で結果を動的に探索できます。数十件の run を一度に簡単に比較し、興味深い結果にズームインして確認したり、柔軟かつインタラクティブな可視化によって複雑なデータからインサイトを引き出したりできます。

上級機能とよくある質問

ベストモデルを保存するにはどうすればよいですか?

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' を設定していた場合は、TrainingArgumentsmodel_name_or_path 引数に model_dir を指定し、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 に接続する
    my_checkpoint_name = f"checkpoint-{last_run_id}:latest"
    my_checkpoint_artifact = run.use_artifact(my_model_name)

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

    # モデルと trainer を再初期化する
    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 ライブラリでは、Transformers の Trainer から W&B へのログ記録は WandbCallback によって処理されます。Hugging Face でのログ出力をカスタマイズする必要がある場合は、WandbCallback をサブクラス化し、Trainer クラスの追加メソッドを活用する機能を追加することで、このコールバックを拡張できます。 以下に、この新しいコールバックを HF Trainer に追加するための一般的なパターンを示し、その後に評価出力を W&B Table にログするコード全体を含む例を示します。
# 通常通りTrainerをインスタンス化する
trainer = Trainer()

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

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

# 通常通りTrainerの学習を開始する
trainer.train()

学習中に評価サンプルを表示する

次のセクションでは、WandbCallback をカスタマイズして、モデルの予測を実行し、評価サンプルを学習中に W&B Table にログする方法を説明します。Trainer のコールバックである on_evaluate メソッドを使って、eval_steps ごとにこれを実行します。 ここでは、トークナイザを使ってモデル出力から予測値とラベルをデコードするために、decode_predictions 関数を作成しています。 次に、予測値とラベルから pandas の DataFrame を作成し、DataFrame に epoch 列を追加します。 最後に、その DataFrame から wandb.Table を作成し、それを wandb にログします。 さらに、freq エポックごとに予測をログすることで、ログの頻度を制御できます。 Note: 通常の WandbCallback と異なり、このカスタムコールバックは Trainer の初期化時ではなく、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):
        """WandbPredictionProgressCallback インスタンスを初期化します。

        Args:
            trainer (Trainer): Hugging Face の Trainer インスタンス。
            tokenizer (AutoTokenizer): モデルに関連付けられた
              トークナイザー。
            val_dataset (Dataset): 検証データセット。
            num_samples (int, optional): 予測生成のために検証データセットから
              選択するサンプル数。
              デフォルトは 100。
            freq (int, optional): ログの頻度。デフォルトは 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 の環境変数の一覧はこちらにあります。
Environment VariableUsage
WANDB_PROJECTプロジェクトに名前を付けます(デフォルトは huggingface
WANDB_LOG_MODEL

モデルのチェックポイントを W&B アーティファクトとしてログに記録します(デフォルトは false

  • false(デフォルト): モデルのチェックポイントは記録しません
  • checkpoint: Trainer の TrainingArguments で設定した args.save_steps ごとにチェックポイントをアップロードします
  • end: 学習の最後に最終モデルのチェックポイントをアップロードします
WANDB_WATCH

モデルの勾配・パラメータを記録するか、あるいはどちらも記録しないかを設定します

  • false(デフォルト): 勾配・パラメータは記録しません
  • gradients: 勾配のヒストグラムを記録します
  • all: 勾配とパラメータのヒストグラムを記録します
WANDB_DISABLEDtrue に設定するとログ記録を完全に無効化します(デフォルトは false
WANDB_QUIET.true に設定すると、標準出力に記録されるメッセージを重大なメッセージのみに制限します(デフォルトは false
WANDB_SILENTtrue に設定すると、wandb が出力するログを抑制します(デフォルトは 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 を作成してください。