メインコンテンツへスキップ
Try in Colab

このノートブックで学べること

  • 実験管理のために、TensorFlow パイプラインを W&B に簡単に統合する方法。
  • keras.metrics を使用したメトリクスの計算。
  • カスタムトレーニングループ内で wandb.log を使用してメトリクスをログ記録する方法。
dashboard
注意: Step で始まるセクションは、既存のコードに W&B を統合するために必要なすべてのステップです。それ以外の部分は標準的な MNIST の例です。
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.datasets import cifar10

インストール、インポート、ログイン

W&B のインストール

%%capture
!pip install wandb

W&B のインポートとログイン

import wandb
from wandb.integration.keras import WandbMetricsLogger

wandb.login()
補足: W&B を初めて使用する場合やログインしていない場合、wandb.login() を実行した後に表示されるリンクからサインアップ/ログインページに移動します。サインアップはワンクリックで簡単に行えます。

データセットの準備

# トレーニングデータセットの準備
BATCH_SIZE = 64
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = np.reshape(x_train, (-1, 784))
x_test = np.reshape(x_test, (-1, 784))

# tf.data を使用して入力パイプラインを構築
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(BATCH_SIZE)

val_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
val_dataset = val_dataset.batch(BATCH_SIZE)

モデルとトレーニングループの定義

def make_model():
    inputs = keras.Input(shape=(784,), name="digits")
    x1 = keras.layers.Dense(64, activation="relu")(inputs)
    x2 = keras.layers.Dense(64, activation="relu")(x1)
    outputs = keras.layers.Dense(10, name="predictions")(x2)

    return keras.Model(inputs=inputs, outputs=outputs)
def train_step(x, y, model, optimizer, loss_fn, train_acc_metric):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)

    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))

    train_acc_metric.update_state(y, logits)

    return loss_value
def test_step(x, y, model, loss_fn, val_acc_metric):
    val_logits = model(x, training=False)
    loss_value = loss_fn(y, val_logits)
    val_acc_metric.update_state(y, val_logits)

    return loss_value

トレーニングループへの wandb.log の追加

def train(
    train_dataset,
    val_dataset,
    model,
    optimizer,
    train_acc_metric,
    val_acc_metric,
    epochs=10,
    log_step=200,
    val_log_step=50,
):
    run = wandb.init(
        project="my-tf-integration",
        config={
            "epochs": epochs,
            "log_step": log_step,
            "val_log_step": val_log_step,
            "architecture": "MLP",
            "dataset": "MNIST",
        },
    )
    for epoch in range(epochs):
        print("\nStart of epoch %d" % (epoch,))

        train_loss = []
        val_loss = []

        # データセットのバッチを反復処理
        for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
            loss_value = train_step(
                x_batch_train,
                y_batch_train,
                model,
                optimizer,
                loss_fn,
                train_acc_metric,
            )
            train_loss.append(float(loss_value))

        # 各エポックの最後に検証ループを実行
        for step, (x_batch_val, y_batch_val) in enumerate(val_dataset):
            val_loss_value = test_step(
                x_batch_val, y_batch_val, model, loss_fn, val_acc_metric
            )
            val_loss.append(float(val_loss_value))

        # 各エポックの最後にメトリクスを表示
        train_acc = train_acc_metric.result()
        print("Training acc over epoch: %.4f" % (float(train_acc),))

        val_acc = val_acc_metric.result()
        print("Validation acc: %.4f" % (float(val_acc),))

        # 各エポックの最後にメトリクスをリセット
        train_acc_metric.reset_state()
        val_acc_metric.reset_state()

        # run.log() を使用してメトリクスをログ記録
        run.log(
            {
                "epochs": epoch,
                "loss": np.mean(train_loss),
                "acc": float(train_acc),
                "val_loss": np.mean(val_loss),
                "val_acc": float(val_acc),
            }
        )
    run.finish()

トレーニングの実行

wandb.init() を呼び出して Run を開始する

これにより、実験の開始が通知され、一意の ID とダッシュボードが割り当てられます。 公式ドキュメントを確認する
# プロジェクト名とオプションの設定(config)を指定して wandb を初期化します。
# config の値を変更して、wandb ダッシュボードでの結果の変化を確認してください。
config = {
    "learning_rate": 0.001,
    "epochs": 10,
    "batch_size": 64,
    "log_step": 200,
    "val_log_step": 50,
    "architecture": "CNN",
    "dataset": "CIFAR-10",
}

run = wandb.init(project='my-tf-integration', config=config)
config = run.config

# モデルの初期化
model = make_model()

# モデルをトレーニングするためのオプティマイザーをインスタンス化
optimizer = keras.optimizers.SGD(learning_rate=config.learning_rate)
# 損失関数をインスタンス化
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# メトリクスの準備
train_acc_metric = keras.metrics.SparseCategoricalAccuracy()
val_acc_metric = keras.metrics.SparseCategoricalAccuracy()

train(
    train_dataset,
    val_dataset, 
    model,
    optimizer,
    train_acc_metric,
    val_acc_metric,
    epochs=config.epochs, 
    log_step=config.log_step, 
    val_log_step=config.val_log_step,
)

run.finish()  # Jupyter/Colab の場合、終了したことを通知します

結果の可視化

上記の Run ページ へのリンクをクリックして、リアルタイムの結果を確認してください。

Sweep 101

W&B Sweeps を使用してハイパーパラメーター最適化を自動化し、可能なモデルの空間を探索します。 W&B Sweeps を使用したハイパーパラメーター最適化の Colab ノートブック をご覧ください。

W&B Sweeps を使用するメリット

  • 迅速なセットアップ: わずか数行のコードで W&B Sweeps を実行できます。
  • 透明性: 使用しているすべてのアルゴリズムを引用しており、コードはオープンソース です。
  • 強力: Sweeps は完全にカスタマイズおよび設定可能です。数十台ののマシンで Sweep を起動するのも、ノート PC で開始するのと同じくらい簡単です。
Sweep result

実例ギャラリー

W&B で追跡および可視化されたプロジェクトの例を、ギャラリー Fully Connected → で探索してください。

ベストプラクティス

  1. Projects: 複数の Runs をプロジェクトにログ記録して比較します。 wandb.init(project="project-name")
  2. Groups: 複数のプロセスや交差検証のフォールドがある場合、各プロセスを Run としてログ記録し、それらをグループ化します。 wandb.init(group="experiment-1")
  3. Tags: タグを追加して、現在のベースラインやプロダクションモデルを追跡します。
  4. Notes: テーブルにノートを入力して、Runs 間の変更を追跡します。
  5. Reports: 進捗に関する簡単なメモを取り、同僚と共有したり、ML プロジェクトのダッシュボードやスナップショットを作成したりします。

高度な設定

  1. 環境変数: 環境変数に API キーを設定して、管理されたクラスター上でトレーニングを実行できるようにします。
  2. オフラインモード
  3. オンプレミス: 独自のインフラストラクチャー内のプライベートクラウドやエアギャップされたサーバーに W&B をインストールします。アカデミックからエンタープライズチームまで、あらゆるユーザー向けのローカルインストールを提供しています。
  4. Artifacts: モデルのトレーニング時にパイプラインのステップを自動的に取得し、モデルやデータセットを合理的な方法で追跡およびバージョン管理します。