Keras

Keras コールバック

W&B は Keras 用に3つのコールバックを提供しています。wandb v0.13.4から利用可能です。レガシーな WandbCallback は下にスクロールしてください。

• WandbMetricsLogger : このコールバックは Experiment Tracking に使用します。トレーニングと検証のメトリクス、システムメトリクスを Weights and Biases にログします。

• WandbModelCheckpoint : モデルのチェックポイントを Weights and Biases の Artifacts にログするためにこのコールバックを使用します。

• WandbEvalCallback: このベースコールバックは、モデルの予測を Weights and Biases の Tables にログして、インタラクティブな可視化を行います。

これらの新しいコールバックは以下の特徴を持っています：

• Keras のデザイン哲学に従います。
• すべての機能に対して単一のコールバック (WandbCallback) を使用する際の認知負荷を減らします。
• Keras ユーザーがコールバックをサブクラス化してニッチなユースケースをサポートできるように簡単に改修できます。

WandbMetricsLogger を使用して実験を追跡

WandbMetricsLogger は、on_epoch_end や on_batch_end などのコールバックメソッドが引数として取得する Keras の logs 辞書を自動的にログします。

これにより次の項目が追跡されます：

• model.compile に定義されたトレーニングと検証のメトリクス。
• システム (CPU/GPU/TPU) のメトリクス。
• 学習率（固定値または学習率スケジューラ）。

import wandb
from wandb.integration.keras import WandbMetricsLogger

# 新しい W&B run を初期化
wandb.init(config={"bs": 12})

# WandbMetricsLogger を model.fit に渡す
model.fit(
    X_train, y_train, validation_data=(X_test, y_test), callbacks=[WandbMetricsLogger()]
)

WandbMetricsLogger リファレンス

WandbModelCheckpoint を使用してモデルをチェックポイント

WandbModelCheckpoint コールバックを使用して、Keras モデル (SavedModel 形式) またはモデルの重みを定期的に保存し、モデルのバージョン管理のために W&B アーティファクトとしてアップロードします。

このコールバックは tf.keras.callbacks.ModelCheckpoint からサブクラス化されているため、チェックポイントのロジックは親コールバックによって処理されます。

このコールバックが保存するもの：

• モニターに基づいて最高のパフォーマンスを達成したモデル。
• パフォーマンスに関係なく各エポック終了時のモデル。
• エポックまたは一定のトレーニングバッチ数後のモデル。
• モデルの重みのみ、またはモデル全体。
• SavedModel 形式または .h5 形式いずれかのモデル。

このコールバックは WandbMetricsLogger と併用してください。

import wandb
from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint

# 新しい W&B run を初期化
wandb.init(config={"bs": 12})

# WandbModelCheckpoint を model.fit に渡す
model.fit(
    X_train,
    y_train,
    validation_data=(X_test, y_test),
    callbacks=[
        WandbMetricsLogger(),
        WandbModelCheckpoint("models"),
    ],
)

WandbModelCheckpoint リファレンス

N エポック後にチェックポイントをログ

デフォルト (save_freq="epoch") では、コールバックは各エポック後にアーティファクトとしてチェックポイントを作成し、アップロードします。特定のバッチ数後にチェックポイントを作成するには、save_freq を整数に設定します。N エポック後にチェックポイントを作成するには、train データローダーの基数を計算し、それを save_freq に渡します。

WandbModelCheckpoint(
    filepath="models/",
    save_freq=int((trainloader.cardinality()*N).numpy())
)

TPU アーキテクチャーで効率的にチェックポイントをログ

TPU 上でチェックポイントを作成する際に、UnimplementedError: File system scheme '[local]' not implemented エラーメッセージが発生することがあります。これは、モデルディレクトリー (filepath) がクラウドストレージバケットパス (gs://bucket-name/...) を使用しなければならないためであり、このバケットは TPU サーバーからアクセス可能でなければなりません。ただし、ローカルパスを使用してチェックポイントを行い、それを Artifacts としてアップロードすることは可能です。

checkpoint_options = tf.saved_model.SaveOptions(experimental_io_device="/job:localhost")

WandbModelCheckpoint(
    filepath="models/,
    options=checkpoint_options,
)

モデル予測を WandbEvalCallback で可視化

WandbEvalCallback は、モデル予測のための Keras コールバックを主に構築するための抽象基底クラスであり、副次的にデータセットの可視化にも使われます。

この抽象コールバックは、データセットやタスクに対してはアグノスティックです。これを使用するには、このベース WandbEvalCallback コールバッククラスを継承し、add_ground_truth と add_model_prediction メソッドを実装します。

WandbEvalCallback は、以下のメソッドを提供するユーティリティクラスです：

• データと予測の wandb.Table インスタンスを作成します。
• データと予測のテーブルを wandb.Artifact としてログします。
• on_train_begin 時にデータテーブルをログします。
• on_epoch_end 時に予測テーブルをログします。

以下の例では、画像分類タスクのために WandbClfEvalCallback を使用しています。この例のコールバックは検証データ (data_table) を W&B にログし、推論を行い、各エポック終了時に予測 (pred_table) を W&B にログします。

import wandb
from wandb.integration.keras import WandbMetricsLogger, WandbEvalCallback


# モデル予測可視化用コールバックを実装
class WandbClfEvalCallback(WandbEvalCallback):
    def __init__(
        self, validation_data, data_table_columns, pred_table_columns, num_samples=100
    ):
        super().__init__(data_table_columns, pred_table_columns)

        self.x = validation_data[0]
        self.y = validation_data[1]

    def add_ground_truth(self, logs=None):
        for idx, (image, label) in enumerate(zip(self.x, self.y)):
            self.data_table.add_data(idx, wandb.Image(image), label)

    def add_model_predictions(self, epoch, logs=None):
        preds = self.model.predict(self.x, verbose=0)
        preds = tf.argmax(preds, axis=-1)

        table_idxs = self.data_table_ref.get_index()

        for idx in table_idxs:
            pred = preds[idx]
            self.pred_table.add_data(
                epoch,
                self.data_table_ref.data[idx][0],
                self.data_table_ref.data[idx][1],
                self.data_table_ref.data[idx][2],
                pred,
            )


# ...

# 新しい W&B run を初期化
wandb.init(config={"hyper": "parameter"})

# コールバックを Model.fit に追加
model.fit(
    X_train,
    y_train,
    validation_data=(X_test, y_test),
    callbacks=[
        WandbMetricsLogger(),
        WandbClfEvalCallback(
            validation_data=(X_test, y_test),
            data_table_columns=["idx", "image", "label"],
            pred_table_columns=["epoch", "idx", "image", "label", "pred"],
        ),
    ],
)

WandbEvalCallback リファレンス

メモリ使用量の詳細

data_table は on_train_begin メソッドが呼び出されたときに W&B にログされます。一度 W&B アーティファクトとしてアップロードされると、data_table_ref クラス変数を使用してこのテーブルにアクセスすることができます。data_table_ref は 2D リストで、self.data_table_ref[idx][n] のようにインデックスを付けてアクセスできます。この例では、idx は行番号で、n は列番号です。

コールバックのカスタマイズ

on_train_begin や on_epoch_end メソッドをオーバーライドして、より細かい制御を行うことができます。N バッチ後にサンプルをログしたい場合、on_train_batch_end メソッドを実装することができます。

WandbCallback [レガシー]

WandbCallback クラスを使用して、model.fit で追跡されるすべてのメトリクスと損失値を自動的に保存します。

import wandb
from wandb.integration.keras import WandbCallback

wandb.init(config={"hyper": "parameter"})

...  # Keras でモデルをセットアップするためのコード

# コールバックを model.fit に渡す
model.fit(
    X_train, y_train, validation_data=(X_test, y_test), callbacks=[WandbCallback()]
)

短いビデオ Get Started with Keras and Weights & Biases in Less Than a Minute をご覧ください。

より詳細なビデオは Integrate Weights & Biases with Keras をご覧ください。Colab Jupyter Notebook を確認できます。

WandbCallback クラスは、広範なロギング設定オプションをサポートしています：監視するメトリクスの指定、重みや勾配の追跡、トレーニングデータと検証データ上の予測のログなど。

keras.WandbCallback の参考文献のドキュメントも確認してください。より詳細な情報があります。

WandbCallback

• Keras によって収集された任意のメトリクスの履歴データを自動的にログします：keras_model.compile() に渡された損失とその他の項目。
• monitor と mode 属性によって定義された “最良” のトレーニングステップに関連付けられたサマリーメトリクスを設定します。これはデフォルトでは最小の val_loss を持つエポックとなります。WandbCallback はデフォルトで最も良い epoch に関連付けられたモデルを保存します。
• 勾配とパラメータのヒストグラムをオプションでログします。
• オプションで wandb に視覚化するためのトレーニングおよび検証データを保存します。

WandbCallback リファレンス

よくある質問

Keras のマルチプロセッシングを wandb でどのように使用しますか？

use_multiprocessing=True を設定したときにこのエラーが発生する場合があります：

Error("You must call wandb.init() before wandb.config.batch_size")

それを回避する方法：

1. Sequence クラスの構築時に、これを追加：wandb.init(group='...')。
2. main で、if __name__ == "__main__": を使用して、スクリプトロジックの残りをその中に置くようにしてください。