メインコンテンツへスキップ
Try in Colab W&B を使用して、機械学習の 実験管理 、データセットの バージョン管理 、プロジェクトのコラボレーションを行いましょう。
W&Bを使用するメリット
この Colabノートブック では、モデルの 予測 の 可視化 や データセット の 可視化 に役立つ コールバック を構築するために継承できる抽象 コールバック である WandbEvalCallback を紹介します。

​
セットアップとインストール

まず、最新バージョンの W&B をインストールします。その後、この Colab インスタンスを W&B で使用できるように認証します。 
pip install -qq -U wandb

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import models
import tensorflow_datasets as tfds

# W&B 関連のインポート
import wandb
from wandb.integration.keras import WandbMetricsLogger
from wandb.integration.keras import WandbModelCheckpoint
from wandb.integration.keras import WandbEvalCallback
W&B を初めて使用する場合、またはログインしていない場合は、 wandb.login() を実行した後に表示されるリンクからサインアップ/ログインページに移動します。無料アカウント への登録は、数クリックで簡単に行えます。 
wandb.login()

​
ハイパーパラメーター

再現可能な 機械学習 のためには、適切な設定（config）システムの使用が推奨されるベストプラクティスです。W&B を使用して、すべての 実験 の ハイパーパラメーター を追跡できます。この Colab では、シンプルな Python の dict を設定システムとして使用します。 
configs = dict(
    num_classes=10,
    shuffle_buffer=1024,
    batch_size=64,
    image_size=28,
    image_channels=1,
    earlystopping_patience=3,
    learning_rate=1e-3,
    epochs=10,
)

​
データセット

この Colab では、TensorFlow Dataset カタログの Fashion-MNIST データセット を使用します。TensorFlow/Keras を使用して、シンプルな 画像分類 パイプライン を構築することを目指します。 
train_ds, valid_ds = tfds.load("fashion_mnist", split=["train", "test"])

AUTOTUNE = tf.data.AUTOTUNE


def parse_data(example):
    # 画像の取得
    image = example["image"]
    # image = tf.image.convert_image_dtype(image, dtype=tf.float32)

    # ラベルの取得
    label = example["label"]
    label = tf.one_hot(label, depth=configs["num_classes"])

    return image, label


def get_dataloader(ds, configs, dataloader_type="train"):
    dataloader = ds.map(parse_data, num_parallel_calls=AUTOTUNE)

    if dataloader_type=="train":
        dataloader = dataloader.shuffle(configs["shuffle_buffer"])
      
    dataloader = (
        dataloader
        .batch(configs["batch_size"])
        .prefetch(AUTOTUNE)
    )

    return dataloader

trainloader = get_dataloader(train_ds, configs)
validloader = get_dataloader(valid_ds, configs, dataloader_type="valid")

​
モデル

def get_model(configs):
    backbone = tf.keras.applications.mobilenet_v2.MobileNetV2(
        weights="imagenet", include_top=False
    )
    backbone.trainable = False

    inputs = layers.Input(
        shape=(configs["image_size"], configs["image_size"], configs["image_channels"])
    )
    resize = layers.Resizing(32, 32)(inputs)
    neck = layers.Conv2D(3, (3, 3), padding="same")(resize)
    preprocess_input = tf.keras.applications.mobilenet.preprocess_input(neck)
    x = backbone(preprocess_input)
    x = layers.GlobalAveragePooling2D()(x)
    outputs = layers.Dense(configs["num_classes"], activation="softmax")(x)

    return models.Model(inputs=inputs, outputs=outputs)

tf.keras.backend.clear_session()
model = get_model(configs)
model.summary()

​
モデルのコンパイル

model.compile(
    optimizer="adam",
    loss="categorical_crossentropy",
    metrics=[
        "accuracy",
        tf.keras.metrics.TopKCategoricalAccuracy(k=5, name="top@5_accuracy"),
    ],
)

​
WandbEvalCallback

WandbEvalCallback は、主にモデルの 予測 の 可視化 、および副次的に データセット の 可視化 のための Keras コールバック を構築するための抽象基底クラスです。 これは データセット やタスクに依存しない抽象 コールバック です。これを使用するには、この基底 コールバック クラスを継承し、 add_ground_truth メソッドと add_model_prediction メソッドを実装します。 WandbEvalCallback は、以下のことを行うのに役立つ メソッド を提供するユーティリティクラスです。
  • データおよび 予測 の wandb.Table インスタンスを作成する。
  • データおよび 予測 の Tables を Artifacts として ログ 記録する。
  • on_train_begin 時に データテーブル を ログ 記録する。
  • on_epoch_end 時に 予測テーブル を ログ 記録する。
例として、以下に 画像分類 タスクのための WandbClfEvalCallback を実装しました。この例の コールバック は以下のことを行います。
  • 検証 データ (data_table) を W&B に ログ 記録する。
  • 推論を実行し、各 エポック 終了時に 予測 (pred_table) を W&B に ログ 記録する。

​
メモリ使用量の削減方法

on_train_begin メソッドが呼び出されたときに data_table を W&B に ログ 記録します。W&B Artifact としてアップロードされると、この テーブル への参照が取得され、 data_table_ref クラス変数を使用して アクセス できるようになります。 data_table_ref は 2次元リストであり、 self.data_table_ref[idx][n] のようにインデックスを指定できます（idx は行番号、 n は列番号）。以下の例で使用方法を確認してください。 
class WandbClfEvalCallback(WandbEvalCallback):
    def __init__(
        self, validloader, data_table_columns, pred_table_columns, num_samples=100
    ):
        super().__init__(data_table_columns, pred_table_columns)

        self.val_data = validloader.unbatch().take(num_samples)

    def add_ground_truth(self, logs=None):
        for idx, (image, label) in enumerate(self.val_data):
            # idx, 画像, 正解ラベルをデータテーブルに追加
            self.data_table.add_data(idx, wandb.Image(image), np.argmax(label, axis=-1))

    def add_model_predictions(self, epoch, logs=None):
        # 予測の取得
        preds = self._inference()
        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,
            )

    def _inference(self):
        preds = []
        for image, label in self.val_data:
            pred = self.model(tf.expand_dims(image, axis=0))
            argmax_pred = tf.argmax(pred, axis=-1).numpy()[0]
            preds.append(argmax_pred)

        return preds

​
学習

# W&B Run を初期化
run = wandb.init(project="intro-keras", config=configs)

# モデルの学習
model.fit(
    trainloader,
    epochs=configs["epochs"],
    validation_data=validloader,
    callbacks=[
        WandbMetricsLogger(log_freq=10),
        WandbClfEvalCallback(
            validloader,
            data_table_columns=["idx", "image", "ground_truth"],
            pred_table_columns=["epoch", "idx", "image", "ground_truth", "prediction"],
        ),  # ここで WandbEvalCallback を使用していることに注目してください
    ],
)

# W&B Run を終了
run.finish()