Tutorial: Use W&B for model management

import wandb
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers
from wandb.integration.keras import WandbMetricsLogger
from sklearn.model_selection import train_test_split

entity = "<entity>"

def generate_raw_data(train_size=6000):
    eval_size = int(train_size / 6)
    (x_train, y_train), (x_eval, y_eval) = keras.datasets.mnist.load_data()

    x_train = x_train.astype("float32") / 255
    x_eval = x_eval.astype("float32") / 255
    x_train = np.expand_dims(x_train, -1)
    x_eval = np.expand_dims(x_eval, -1)

    print("Generated {} rows of training data.".format(train_size))
    print("Generated {} rows of eval data.".format(eval_size))

    return (x_train[:train_size], y_train[:train_size]), (
        x_eval[:eval_size],
        y_eval[:eval_size],
    )

# データセットを作成します
(x_train, y_train), (x_eval, y_eval) = generate_raw_data()

project = "model-registry-dev"

model_use_case_id = "mnist"
job_type = "build_dataset"

# W&B の run を初期化します
run = wandb.init(entity=entity, project=project, job_type=job_type)

# トレーニングデータ用の W&B テーブルを作成します
train_table = wandb.Table(data=[], columns=[])
train_table.add_column("x_train", x_train)
train_table.add_column("y_train", y_train)
train_table.add_computed_columns(lambda ndx, row: {"img": wandb.Image(row["x_train"])})

# 評価データ用の W&B テーブルを作成します
eval_table = wandb.Table(data=[], columns=[])
eval_table.add_column("x_eval", x_eval)
eval_table.add_column("y_eval", y_eval)
eval_table.add_computed_columns(lambda ndx, row: {"img": wandb.Image(row["x_eval"])})

# Artifact オブジェクトを作成します
artifact_name = "{}_dataset".format(model_use_case_id)
artifact = wandb.Artifact(name=artifact_name, type="dataset")

# wandb.WBValue オブジェクトを Artifact に追加します。
artifact.add(train_table, "train_table")
artifact.add(eval_table, "eval_table")

# Artifact に加えられた変更を永続化します。
artifact.save()

# W&B にこの run が完了したことを伝えます。
run.finish()

job_type = "train_model"
config = {
    "optimizer": "adam",
    "batch_size": 128,
    "epochs": 5,
    "validation_split": 0.1,
}

# W&B の run を初期化します
run = wandb.init(project=project, job_type=job_type, config=config)

# データセット Artifact を取得します
version = "latest"
name = "{}:{}".format("{}_dataset".format(model_use_case_id), version)
artifact = run.use_artifact(artifact_or_name=name)

# データフレームから特定の内容を取得します
train_table = artifact.get("train_table")
x_train = train_table.get_column("x_train", convert_to="numpy")
y_train = train_table.get_column("y_train", convert_to="numpy")

path = "./model.h5"
registered_model_name = "MNIST-dev"

run.link_model(path=path, registered_model_name=registered_model_name)
run.finish()

job_type = "evaluate_model"

# run を初期化します
run = wandb.init(project=project, entity=entity, job_type=job_type)

model_use_case_id = "mnist"
version = "latest"

# データセット Artifact を取得し、依存関係としてマークします
artifact = run.use_artifact(
    "{}:{}".format("{}_dataset".format(model_use_case_id), version)
)

# 目的のデータフレームを取得します
eval_table = artifact.get("eval_table")
x_eval = eval_table.get_column("x_eval", convert_to="numpy")
y_eval = eval_table.get_column("y_eval", convert_to="numpy")

alias = "latest"  # エイリアス
name = "mnist_model"  # モデル Artifact の名前

# モデルにアクセスしてダウンロードします。ダウンロードした Artifact へのパスを返します
downloaded_model_path = run.use_model(name=f"{name}:{alias}")

model = keras.models.load_model(downloaded_model_path)

y_eval = keras.utils.to_categorical(y_eval, 10)
(loss, _) = model.evaluate(x_eval, y_eval)
score = (loss, _)

# # メトリクス、画像、テーブル、または評価に役立つデータを記録します。
run.log(data={"loss": (loss, _)})