Track models and datasets

Try in a Colab Notebook here →

このノートブックでは、W&B Artifacts を使ってMLの実験パイプラインをトラッキングする方法を紹介します。

一緒に進めましょう！ ビデオチュートリアル

🤔 Artifactsとは何か、そしてなぜ重要なのか？

「artifact」は、ギリシャのアンフォラ 🏺のような作られたオブジェクトで、プロセスの出力物です。 MLでは、最も重要なartifactは datasets と models です。

そして、コロナドの十字架のように、これらの重要なartifactは博物館に保存されるべきです！ つまり、これらはカタログ化され、整理されるべきです。 あなたやチーム、そして広範なMLコミュニティがそれらから学べるようにするためです。 トレーニングをトラッキングしない者はそれを繰り返す運命にあるのです。

Artifacts APIを使用すると、ArtifactをW&BのRunの出力としてログに記録したり、Runの入力としてArtifactを使用したりできます。このダイアグラムのように、トレーニングrunがデータセットを入力として取り、モデルを生成します。

1つのrunが他のrunの出力を入力として使用できるため、ArtifactsとRunsは一緒に有向グラフ（実際には二部ダイアグラム）を形成します。ノードはArtifactとRunのものであり、それぞれ消費や生成されるArtifactに対してRunsが矢印でつながっています。

0️⃣ インストールとインポート

Artifactsは、バージョン0.9.2からPythonライブラリの一部です。

他の多くのML Pythonスタックと同様に、pip経由で利用可能です。

# wandb バージョン0.9.2以降に対応
!pip install wandb -qqq
!apt install tree

import os
import wandb

1️⃣ データセットをログする

まず、いくつかのArtifactsを定義しましょう。

この例はPyTorchの "Basic MNIST Example" に基づいていますが、TensorFlowや他の任意のフレームワーク、または純粋なPythonでも簡単に実行できます。

まずは、Datasetを以下のように定義します：

• trainセットはパラメータを選ぶため、
• validationセットはハイパーパラメータを選ぶため、
• testセットは最終モデルを評価するため

以下のセルでは、これら3つのデータセットを定義します。

import random 

import torch
import torchvision
from torch.utils.data import TensorDataset
from tqdm.auto import tqdm

# 決定的な振る舞いを保証
torch.backends.cudnn.deterministic = True
random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

# デバイスの設定
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# データパラメータ
num_classes = 10
input_shape = (1, 28, 28)

# 遅いミラーをドロップ
torchvision.datasets.MNIST.mirrors = [mirror for mirror in torchvision.datasets.MNIST.mirrors
                                      if not mirror.startswith("http://yann.lecun.com")]

def load(train_size=50_000):
    """
    # データをロード
    """

    # データをトレーニングおよびテストセットに分割
    train = torchvision.datasets.MNIST("./", train=True, download=True)
    test = torchvision.datasets.MNIST("./", train=False, download=True)
    (x_train, y_train), (x_test, y_test) = (train.data, train.targets), (test.data, test.targets)

    # ハイパーパラメータチューニングのために検証セットを分割
    x_train, x_val = x_train[:train_size], x_train[train_size:]
    y_train, y_val = y_train[:train_size], y_train[train_size:]

    training_set = TensorDataset(x_train, y_train)
    validation_set = TensorDataset(x_val, y_val)
    test_set = TensorDataset(x_test, y_test)

    datasets = [training_set, validation_set, test_set]

    return datasets

このコードは、データを生成するコードの周りにデータをArtifactとしてログするコードがラップされています。 この場合、データをloadするコードと、データをload_and_logするコードが分離されています。

これは良い実践です！

これらのデータセットをArtifactとしてログするには、

1. wandb.initでRunを作成 (L4)
2. データセット用のArtifactを作成 (L10)
3. 関連するfileを保存してログ (L20, L23)

以下のコードセルの例を見て、その後のセクションを展開して詳細を確認してください。

def load_and_log():

    # 🚀 runをスタートし、タイプとプロジェクトでラベル付けする
    with wandb.init(project="artifacts-example", job_type="load-data") as run:
        
        datasets = load()  # データセットのロードのための分離コード
        names = ["training", "validation", "test"]

        # 🏺 Artifactを作成
        raw_data = wandb.Artifact(
            "mnist-raw", type="dataset",
            description="生のMNISTデータセットをトレーニング/検証/テストに分割",
            metadata={"source": "torchvision.datasets.MNIST",
                      "sizes": [len(dataset) for dataset in datasets]})

        for name, data in zip(names, datasets):
            # 🐣 アーティファクトに新しいファイルを保存し、その内容に何かを書き込みます。
            with raw_data.new_file(name + ".pt", mode="wb") as file:
                x, y = data.tensors
                torch.save((x, y), file)

        # ✍️ アーティファクトをW&Bに保存します。
        run.log_artifact(raw_data)

load_and_log()

🚀 wandb.init

Artifactを生成するRunを作成する際、どのprojectに属するかを指定する必要があります。

ワークフローに応じて、プロジェクトはcar-that-drives-itselfのように大きいものからiterative-architecture-experiment-117のように小さいものまで様々です。

Rule of 👍: すべてのRunが共有するArtifactを1つのプロジェクト内に保持できる場合は、それをシンプルに保ちましょう。ただし、心配しないでください──Artifactはプロジェクト間で移動可能です！

様々なジョブを管理するために、Runを作成する際にjob_typeを提供するのが便利です。 これによりArtifactのグラフが整理されます。

Rule of 👍: job_typeは説明的であり、パイプラインの単一ステップに対応するようにしましょう。ここでは、データのloadとデータのpreprocessを分けています。

🏺 wandb.Artifact

何かをArtifactとしてログするには、最初にArtifactオブジェクトを作成する必要があります。

すべてのArtifactにはnameがあります──これが最初の引数で設定されます。

Rule of 👍: nameは説明的で、覚えやすく入力しやすいものにしましょう──私たちはハイフンで区切られた名前をコード内の変数名に対応させるのが好きです。

また、typeもあります。これは、Runのjob_typeと同様にRunとArtifactのグラフを整理するために使用されます。

Rule of 👍: typeはシンプルであるべきです：datasetやmodelのようなもの。

さらに、辞書形式のdescriptionとmetadataを追加することもできます。 metadataはJSONにシリアライズ可能である必要があります。

Rule of 👍: metadataは可能な限り詳細にします。

🐣 artifact.new_fileと✍️ run.log_artifact

Artifactオブジェクトを作成したら、ファイルを追加する必要があります。

その通り、ファイル の s があります。 Artifactはディレクトリのように構造化されており、ファイルやサブディレクトリを持ちます。

Rule of 👍: できる限り、Artifactの内容を複数のファイルに分割してください。これによりスケールする際に便利です！

new_fileメソッドを使用してファイルを書くと同時にArtifactに添付することができます。 以下では、これらの2つのステップを分けるadd_fileメソッドを使用します。

すべてのファイルを追加したら、log_artifactを使用してwandb.aiに保存する必要があります。

出力にはいくつかのURLが表示されますが、その中にはRunページのURLも含まれます。 そこでは、Runの結果を確認でき、ログされたArtifactも表示されます。

以下はその他のRunページコンポーネントをより活用する例です。

2️⃣ ログされたデータセットArtifactを使用する

W&BのArtifactは、博物館のアーティファクトとは違い、単に保存されるだけでなく、使用されることを想定しています。

それがどのようなものか見てみましょう。

以下のセルでは、生のデータセットを取り込み、それをpreprocessedデータセットとして出力するパイプラインステップを定義します。 これはnormalizeされ、正しく整形されたデータです。

ここでも、wandbとインターフェイスするコードから、重要なコードであるpreprocessを分離しています。

def preprocess(dataset, normalize=True, expand_dims=True):
    """
    ## データを準備する
    """
    x, y = dataset.tensors

    if normalize:
        # 画像を[0, 1]の範囲にスケールする
        x = x.type(torch.float32) / 255

    if expand_dims:
        # 画像が(1, 28, 28)の形状を持っていることを確認する
        x = torch.unsqueeze(x, 1)
    
    return TensorDataset(x, y)

次に、preprocessステップをwandb.Artifactのログを使用して器用に扱うコードです。

以下の例では、Artifactをuseすること（これは新しい）、そしてそれをログすること（これは前のステップと同じ）を行っています。 ArtifactはRunの入力と出力の両方です！

新しいjob_typeであるpreprocess-dataを使用して、これが前のloadとは異なる種類のジョブであることを明確にします。

def preprocess_and_log(steps):

    with wandb.init(project="artifacts-example", job_type="preprocess-data") as run:

        processed_data = wandb.Artifact(
            "mnist-preprocess", type="dataset",
            description="前処理済みのMNISTデータセット",
            metadata=steps)
         
        # ✔️ 使用するartifactを宣言
        raw_data_artifact = run.use_artifact('mnist-raw:latest')

        # 📥 必要に応じてartifactをダウンロード
        raw_dataset = raw_data_artifact.download()
        
        for split in ["training", "validation", "test"]:
            raw_split = read(raw_dataset, split)
            processed_dataset = preprocess(raw_split, **steps)

            with processed_data.new_file(split + ".pt", mode="wb") as file:
                x, y = processed_dataset.tensors
                torch.save((x, y), file)

        run.log_artifact(processed_data)


def read(data_dir, split):
    filename = split + ".pt"
    x, y = torch.load(os.path.join(data_dir, filename))

    return TensorDataset(x, y)

ここで気づくべき点は、前処理のstepsがpreprocessed_dataのmetadataとして保存されることです。

実験を再現可能にしたい場合、たくさんのメタデータをキャプチャすることは良い考えです！

また、私たちのデータセットが「large artifact」であっても、downloadステップは1秒未満で完了します。

詳細は以下のセルを展開して確認してください。

steps = {"normalize": True,
         "expand_dims": True}

preprocess_and_log(steps)

✔️ run.use_artifact

これらのステップはシンプルです。消費者は単にArtifactのnameを知る必要があります、ちょっとだけ他の情報も。

その「ちょっとだけ他の情報」とは、求める特定のArtifactバージョンのaliasです。

デフォルトでは、最後にアップロードされたバージョンにはlatestがタグ付けされます。 それ以外の場合、v0やv1などの古いバージョンを選択することもできますし、 bestやjit-scriptなど自分のエイリアスを提供することもできます。 ちょうどDocker Hubのタグのように、 エイリアスは名前から:で分離されています。

Rule of 👍: エイリアスは短くシンプルに保ちましょう。 カスタムエイリアスlatestやbestを使用して、特定のプロパティを満たすArtifactを選びましょう。

📥 artifact.download

download呼び出しが心配かもしれません。 別のコピーをダウンロードすると、メモリの負担が倍増するのでは？

心配しないでください。実際に何かをダウンロードする前に、 適切なバージョンがローカルに存在するか確認します。 これにはトレントやgitによるバージョン管理の技術が使われています。

Artifactが作成・ログされると、作業ディレクトリの中にartifactsというフォルダができ、 それぞれのArtifact用のサブディレクトリが埋め尽くされます。 以下のコマンドでその内容をチェックしてください：

!tree artifacts

🌐 wandb.aiのArtifactsページ

Artifactをログして使用したら、RunページのArtifactsタブをチェックしましょう。

wandbの出力に表示されるRunページURLに移動し、 左サイドバーから「Artifacts」タブを選択します （これはデータベースアイコンのもので、ホッケーのパックが3つ積み重なったように見えます）。

"Input Artifacts"テーブルや"Output Artifacts"テーブルの行をクリックし、 タブ（"Overview"、"Metadata"）をチェックして、Artifactに関するすべてのログを確認してください。

特に「Graph View」が好きです。 デフォルトでは、ArtifactのタイプとRunのjob_typeが2種類のノードで示され、 矢印は消費と生成を表します。

3️⃣ モデルをログする

これでArtifact APIの動作がわかりましたが、この例をパイプラインの終わりまで続けて、 ArtifactがMLワークフローをどのように改善できるか見てみましょう。

この最初のセルは、PyTorchでDNN modelを構築します──ごくシンプルなConvNetです。

まず、モデルを初期化し、トレーニングはしません。 これにより、他の部分を固定してトレーニングを繰り返すことができます。

from math import floor

import torch.nn as nn

class ConvNet(nn.Module):
    def __init__(self, hidden_layer_sizes=[32, 64],
                  kernel_sizes=[3],
                  activation="ReLU",
                  pool_sizes=[2],
                  dropout=0.5,
                  num_classes=num_classes,
                  input_shape=input_shape):
      
        super(ConvNet, self).__init__()

        self.layer1 = nn.Sequential(
              nn.Conv2d(in_channels=input_shape[0], out_channels=hidden_layer_sizes[0], kernel_size=kernel_sizes[0]),
              getattr(nn, activation)(),
              nn.MaxPool2d(kernel_size=pool_sizes[0])
        )
        self.layer2 = nn.Sequential(
              nn.Conv2d(in_channels=hidden_layer_sizes[0], out_channels=hidden_layer_sizes[-1], kernel_size=kernel_sizes[-1]),
              getattr(nn, activation)(),
              nn.MaxPool2d(kernel_size=pool_sizes[-1])
        )
        self.layer3 = nn.Sequential(
              nn.Flatten(),
              nn.Dropout(dropout)
        )

        fc_input_dims = floor((input_shape[1] - kernel_sizes[0] + 1) / pool_sizes[0]) # layer 1 output size
        fc_input_dims = floor((fc_input_dims - kernel_sizes[-1] + 1) / pool_sizes[-1]) # layer 2 output size
        fc_input_dims = fc_input_dims*fc_input_dims*hidden_layer_sizes[-1] # layer 3 output size

        self.fc = nn.Linear(fc_input_dims, num_classes)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.fc(x)
        return x

ここでは、W&Bを使用してrunをトラッキングし、 wandb.config オブジェクトを使用してすべてのハイパーパラメータを保存しています。

そのconfigオブジェクトのdictバージョンは非常に便利なメタデータなので、必ず含めてください！

def build_model_and_log(config):
    with wandb.init(project="artifacts-example", job_type="initialize", config=config) as run:
        config = wandb.config
        
        model = ConvNet(**config)

        model_artifact = wandb.Artifact(
            "convnet", type="model",
            description="シンプルなAlexNetスタイルのCNN",
            metadata=dict(config))

        torch.save(model.state_dict(), "initialized_model.pth")
        # ➕ Artifactにファイルを追加する別の方法
        model_artifact.add_file("initialized_model.pth")

        wandb.save("initialized_model.pth")

        run.log_artifact(model_artifact)

model_config = {"hidden_layer_sizes": [32, 64],
                "kernel_sizes": [3],
                "activation": "ReLU",
                "pool_sizes": [2],
                "dropout": 0.5,
                "num_classes": 10}

build_model_and_log(model_config)

➕ artifact.add_file

データセットのログ例のように、new_fileを一度に書き込み、Artifactに追加する代わりに、 ファイルを1ステップで書く（ここでは、torch.save）ことができ、その後でArtifactに追加できます。

Rule of 👍: 重複を防ぐために、可能であればnew_fileを使用しましょう。

4️⃣ ログされたモデルArtifactを使用する

データセットに対してuse_artifactを呼び出せるように、 initialized_modelに対しても同様に呼び出して、別のRunで使用できます。

今回はモデルをtrainしましょう。

詳細は、私たちのColabノートブック W&BとPyTorchを使ったインスツルメントをご覧ください。

import torch.nn.functional as F

def train(model, train_loader, valid_loader, config):
    optimizer = getattr(torch.optim, config.optimizer)(model.parameters())
    model.train()
    example_ct = 0
    for epoch in range(config.epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()

            example_ct += len(data)

            if batch_idx % config.batch_log_interval == 0:
                print('Train Epoch: {} [{}/{} ({:.0%})]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                    batch_idx / len(train_loader), loss.item()))
                
                train_log(loss, example_ct, epoch)

        # 各エポックで検証セット上のモデルを評価
        loss, accuracy = test(model, valid_loader)  
        test_log(loss, accuracy, example_ct, epoch)

    
def test(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum')  # バッチ損失を合計
            pred = output.argmax(dim=1, keepdim=True)  # 最大対数確率のインデックスを取得
            correct += pred.eq(target.view_as(pred)).sum()

    test_loss /= len(test_loader.dataset)

    accuracy = 100. * correct / len(test_loader.dataset)
    
    return test_loss, accuracy


def train_log(loss, example_ct, epoch):
    loss = float(loss)

    # マジックが起きる場所
    wandb.log({"epoch": epoch, "train/loss": loss}, step=example_ct)
    print(f"Loss after " + str(example_ct).zfill(5) + f" examples: {loss:.3f}")
    

def test_log(loss, accuracy, example_ct, epoch):
    loss = float(loss)
    accuracy = float(accuracy)

    # マジックが起きる場所
    wandb.log({"epoch": epoch, "validation/loss": loss, "validation/accuracy": accuracy}, step=example_ct)
    print(f"Loss/accuracy after " + str(example_ct).zfill(5) + f" examples: {loss:.3f}/{accuracy:.3f}")

今回は、Artifactを生成する2つの別々のRunを行います。

まず、最初のrunがモデルをトレーニングすると、 次のrunがtrained-modelArtifactを使用して、 テストデータセットでパフォーマンスを評価します。

また、ネットワークが最もうまくいかなかった32例──categorical_crossentropyが最も高い例──を抽出します。

これは、データセットおよびモデルの問題を診断するための良い方法です！

def evaluate(model, test_loader):
    """
    ## 訓練済みモデルを評価
    """

    loss, accuracy = test(model, test_loader)
    highest_losses, hardest_examples, true_labels, predictions = get_hardest_k_examples(model, test_loader.dataset)

    return loss, accuracy, highest_losses, hardest_examples, true_labels, predictions

def get_hardest_k_examples(model, testing_set, k=32):
    model.eval()

    loader = DataLoader(testing_set, 1, shuffle=False)

    # データセット内の各項目の損失と予測を取得
    losses = None
    predictions = None
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = F.cross_entropy(output, target)
            pred = output.argmax(dim=1, keepdim=True)
            
            if losses is None:
                losses = loss.view((1, 1))
                predictions = pred
            else:
                losses = torch.cat((losses, loss.view((1, 1))), 0)
                predictions = torch.cat((predictions, pred), 0)

    argsort_loss = torch.argsort(losses, dim=0)

    highest_k_losses = losses[argsort_loss[-k:]]
    hardest_k_examples = testing_set[argsort_loss[-k:]][0]
    true_labels = testing_set[argsort_loss[-k:]][1]
    predicted_labels = predictions[argsort_loss[-k:]]

    return highest_k_losses, hardest_k_examples, true_labels, predicted_labels

これらのログ関数は、新しいArtifact機能を追加するわけではありません。 単に使用し、ダウンロードし、Artifactをログします。

from torch.utils.data import DataLoader

def train_and_log(config):

    with wandb.init(project="artifacts-example", job_type="train", config=config) as run:
        config = wandb.config

        data = run.use_artifact('mnist-preprocess:latest')
        data_dir = data.download()

        training_dataset =  read(data_dir, "training")
        validation_dataset = read(data_dir, "validation")

        train_loader = DataLoader(training_dataset, batch_size=config.batch_size)
        validation_loader = DataLoader(validation_dataset, batch_size=config.batch_size)
        
        model_artifact = run.use_artifact("convnet:latest")
        model_dir = model_artifact.download()
        model_path = os.path.join(model_dir, "initialized_model.pth")
        model_config = model_artifact.metadata
        config.update(model_config)

        model = ConvNet(**model_config)
        model.load_state_dict(torch.load(model_path))
        model = model.to(device)
 
        train(model, train_loader, validation_loader, config)

        model_artifact = wandb.Artifact(
            "trained-model", type="model",
            description="トレーニング済みNNモデル",
            metadata=dict(model_config))

        torch.save(model.state_dict(), "trained_model.pth")
        model_artifact.add_file("trained_model.pth")
        wandb.save("trained_model.pth")

        run.log_artifact(model_artifact)

    return model

    
def evaluate_and_log(config=None):
    
    with wandb.init(project="artifacts-example", job_type="report", config=config) as run:
        data = run.use_artifact('mnist-preprocess:latest')
        data_dir = data.download()
        testing_set = read(data_dir, "test")

        test_loader = torch.utils.data.DataLoader(testing_set, batch_size=128, shuffle=False)

        model_artifact = run.use_artifact("trained-model:latest")
        model_dir = model_artifact.download()
        model_path = os.path.join(model_dir, "trained_model.pth")
        model_config = model_artifact.metadata

        model = ConvNet(**model_config)
        model.load_state_dict(torch.load(model_path))
        model.to(device)

        loss, accuracy, highest_losses, hardest_examples, true_labels, preds = evaluate(model, test_loader)

        run.summary.update({"loss": loss, "accuracy": accuracy})

        wandb.log({"high-loss-examples":
            [wandb.Image(hard_example, caption=str(int(pred)) + "," +  str(int(label)))
             for hard_example, pred, label in zip(hardest_examples, preds, true_labels)]})

train_config = {"batch_size": 128,
                "epochs": 5,
                "batch_log_interval": 25,
                "optimizer": "Adam"}

model = train_and_log(train_config)
evaluate_and_log()

🔁 グラフビュー

Artifactのtypeを変更しました： これらのRunsはdatasetではなくmodelを使用しました。 グラフビューのArtifactsページでは、 modelを生成するRunsはdatasetを生成するRunsとは区別されます。

ぜひチェックしてみてください！以前と同じように、Runページにアクセスし、 左サイドバーから「Artifacts」タブを選び、 Artifactを選択して「Graph View」タブをクリックしてください。

💣 グラフの展開

「Explode」とラベル付けされたボタンがあるのに気づいたかもしれません。それをクリックしないでください、それは私の机の下に小さな爆弾をセットします…というのは冗談です。そのボタンをクリックすると、グラフがより優しく展開されます： ArtifactやRunsは、個々のインスタンスのレベルで分離されます、タイプではなく、例えば、 ノードはdatasetやload-dataではなく、dataset:mnist-raw:v1やload-data:sunny-smoke-1などです。

これにより、ログされたメトリクス、メタデータなどすべてが手の届く範囲にあり、 パイプラインの全貌を完全に把握できます──あなたが私たちにログを選んだものだけが限界です。

次は何をする？