PyTorch - CoreWeave Docs

W&B を使用して、機械学習の実験管理、データセットのバージョン管理、project でのコラボレーションを行いましょう。

このノートブックの内容

このチュートリアルでは、W&B を PyTorch のトレーニングコードに統合し、実験のトラッキング、メトリクスと勾配のログ、モデルのバージョン管理を行う方法を説明します。既存の PyTorch パイプラインに実験管理を追加する際にご活用ください。

# ライブラリをインポートする
import wandb

# config でハイパーパラメーターの辞書を取得する
config = {
    "learning_rate": 0.001,
    "epochs": 100,
    "batch_size": 128
}

# 新しい実験を開始する
with wandb.init(project="new-sota-model", config=config) as run:

    # モデルとデータを設定する
    model, dataloader = get_model(), get_data()

    # オプション: 勾配をトラッキングする
    run.watch(model)

    for batch in dataloader:
        metrics = model.training_step()
        # トレーニングループ内でメトリクスをログし、モデル性能を可視化する
        run.log(metrics)

    # オプション: 最後にモデルを保存する
    model.to_onnx()
    run.save("model.onnx")

動画チュートリアルを見ながら進めることもできます。既存のパイプラインに W&B を統合するには、Step で始まるセクションだけを実行すれば十分です。それ以外のセクションでは、データの読み込みとモデルの定義を行います。

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

実験を定義する前に、環境を設定して W&B で認証します。

import os
import random

import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from tqdm.auto import tqdm

# 決定論的な動作を保証する
seed = 42
torch.backends.cudnn.deterministic = True
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

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

# MNIST のミラー一覧から低速なミラーを除外する
torchvision.datasets.MNIST.mirrors = [mirror for mirror in torchvision.datasets.MNIST.mirrors
                                      if not mirror.startswith("http://yann.lecun.com")]

ステップ 0: W&B をインストールする

まず、pip を使用して wandb ライブラリをインストールします。

!pip install wandb onnx -Uq

ステップ 1: W&B をインポートしてログインする

W&B サービスにデータをログするには、ログインが必要です。 W&B を初めて使用する場合は、表示されるリンクから無料アカウントに登録してください。

import wandb

wandb.login()

実験とパイプラインを定義する

W&B のインストールとセッションの認証が完了したら、実験の設定と、その設定を使用するトレーニングパイプラインを定義します。

`wandb.init()` でメタデータとハイパーパラメーターをトラッキングする

まず、プログラムで実験を定義します。ハイパーパラメーターには何があるでしょうか。この run にはどのようなメタデータが関連付けられるでしょうか。一般的なワークフローでは、これらの情報を config 辞書 (または同様のオブジェクト) に格納し、必要に応じて参照します。この例では、一部のハイパーパラメーターのみを変化させ、残りはハードコーディングしています。モデルのどの部分でも config に含めることができます。この例には、MNIST データセットと畳み込みアーキテクチャに関するメタデータも含まれています。後で同じ project で CIFAR を使って全結合アーキテクチャを扱う場合なども、このメタデータを使って run を区別できます。

config = dict(
    epochs=5,
    classes=10,
    kernels=[16, 32],
    batch_size=128,
    learning_rate=0.005,
    dataset="MNIST",
    architecture="CNN")

次に、パイプライン全体を定義します。これはモデルのトレーニングでよく使われる構成です。

モデルと、それに関連するデータおよびオプティマイザーを make します。
それに応じてモデルを train します。
test を実行して、トレーニングの結果を確認します。

以下のコードで、これらの関数を実装します。

def model_pipeline(hyperparameters):

    # wandb を開始します
    with wandb.init(project="pytorch-demo", config=hyperparameters) as run:
        # ログする内容と実際の実行内容を一致させるため、すべてのハイパーパラメーターには run.config 経由でアクセスします。
        config = run.config

        # モデル、データ、最適化問題を作成します
        model, train_loader, test_loader, criterion, optimizer = make(config)
        print(model)

        # 作成したものを使用してモデルをトレーニングします
        train(model, train_loader, criterion, optimizer, config)

        # 最終的なパフォーマンスをテストします
        test(model, test_loader)

    return model

標準的なパイプラインとの唯一の違いは、すべての処理が wandb.init() のコンテキスト内で実行される点です。この関数を呼び出すと、コードと W&B サーバーとの間に通信経路が確立されます。 config 辞書を wandb.init() に渡すと、その情報はすべてただちに W&B にログされます。そのため、実験で使用するよう設定したハイパーパラメーターの値を常に把握できます。選択してログした値が確実にモデルで使用されるよう、W&B ではオブジェクトの run.config コピーを使用することを推奨しています。具体例については、以下の make の定義を参照してください。パイプラインを定義したら、以降のセクションでは、データとモデルのセットアップ、トレーニング、テストの各ステップを順に実装していきます。

補足: W&B は独立したプロセスでコードを実行するため、W&B 側で問題が発生してもコードがクラッシュすることはありません。問題が解決したら、wandb sync を使用してデータをログできます。

def make(config):
    # データを準備
    train, test = get_data(train=True), get_data(train=False)
    train_loader = make_loader(train, batch_size=config.batch_size)
    test_loader = make_loader(test, batch_size=config.batch_size)

    # モデルを作成
    model = ConvNet(config.kernels, config.classes).to(device)

    # 損失関数とオプティマイザーを作成
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=config.learning_rate)
    
    return model, train_loader, test_loader, criterion, optimizer

データの読み込みとモデルを定義する

次に、データの読み込み方法とモデルの構成を指定します。この部分は重要ですが、内容は wandb を使わない場合と変わりません。

def get_data(slice=5, train=True):
    full_dataset = torchvision.datasets.MNIST(root=".",
                                              train=train, 
                                              transform=transforms.ToTensor(),
                                              download=True)
    #  [::slice] でスライスするのと同じ 
    sub_dataset = torch.utils.data.Subset(
      full_dataset, indices=range(0, len(full_dataset), slice))
    
    return sub_dataset


def make_loader(dataset, batch_size):
    loader = torch.utils.data.DataLoader(dataset=dataset,
                                         batch_size=batch_size, 
                                         shuffle=True,
                                         pin_memory=True, num_workers=2)
    return loader

wandb を使用してもモデルの定義方法は変わらないため、この例では標準的な ConvNet アーキテクチャを使用します。このコードを自由に変更して実験してください。W&B はすべての結果を Forge にログします。

# 従来型および畳み込みニューラルネットワーク

class ConvNet(nn.Module):
    def __init__(self, kernels, classes=10):
        super(ConvNet, self).__init__()
        
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, kernels[0], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, kernels[1], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7 * 7 * kernels[-1], classes)
        
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

トレーニングロジックを定義する

model_pipeline の次のステップとして、train の方法を指定します。ここでは、トレーニングの進行に合わせて、W&B インテグレーションが勾配、パラメーター、メトリクスをトラッキングします。ここで使用する wandb 関数は、watch と log の 2 つです。

`run.watch()` で勾配をトラッキングし、それ以外はすべて `run.log()` でトラッキングする

run.watch() は、トレーニングの log_freq ステップごとに、モデルの勾配とパラメーターをログします。 run.watch() はトレーニングを開始する前に呼び出してください。ログモード、複数のモデルの扱い、パフォーマンスに関するヒントについては、wandb.watch で勾配とモデルの重みをログするにはどうすればよいですか？を参照してください。それ以外のトレーニングコードは変更不要です。エポックとバッチを反復処理し、順伝播と逆伝播を実行して、optimizer を適用します。

def train(model, loader, criterion, optimizer, config):
    # wandb でモデルの挙動（勾配、重みなど）を監視します。
    run = wandb.init(project="pytorch-demo", config=config)
    run.watch(model, criterion, log="all", log_freq=10)

    # トレーニングを実行し、wandb でトラッキングします
    total_batches = len(loader) * config.epochs
    example_ct = 0  # 処理済みのサンプル数
    batch_ct = 0
    for epoch in tqdm(range(config.epochs)):
        for _, (images, labels) in enumerate(loader):

            loss = train_batch(images, labels, model, optimizer, criterion)
            example_ct +=  len(images)
            batch_ct += 1

            # 25 バッチごとにメトリクスを記録します
            if ((batch_ct + 1) % 25) == 0:
                train_log(loss, example_ct, epoch)


def train_batch(images, labels, model, optimizer, criterion):
    images, labels = images.to(device), labels.to(device)
    
    # 順伝播 ➡
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # 逆伝播 ⬅
    optimizer.zero_grad()
    loss.backward()

    # オプティマイザーでパラメーターを更新します
    optimizer.step()

    return loss

唯一の違いはログするコードの部分です。これまではメトリクスをターミナルに出力して報告していたかもしれませんが、今後は同じ情報を run.log() に渡します。 run.log() は、strings をキーとする辞書を受け取ります。これらの strings は、値として渡されるログ対象のオブジェクトを識別します。また、オプションとして、トレーニングの現在の step をログすることもできます。

補足: モデルが処理したサンプル数を使用すると、異なるバッチサイズ間で比較しやすくなりますが、ステップ数やバッチ数をそのまま使用することもできます。長時間にわたるトレーニング run では、epoch 単位でログするのも有効です。

def train_log(loss, example_ct, epoch):
    with wandb.init(project="pytorch-demo") as run:
        # 損失とエポック番号をログする
        # ここでメトリクスを W&B にログする
        run.log({"epoch": epoch, "loss": loss}, step=example_ct)
        print(f"Loss after {str(example_ct).zfill(5)} examples: {loss:.3f}")

テストロジックを定義する

モデルのトレーニングが完了したら、モデルをテストします。たとえば、本番環境から取得した新しいデータでモデルを実行したり、手作業で厳選したサンプルに適用したりします。テストの段階は、トレーニング済みモデルを保存するのにも適したタイミングです。

オプション: `run.save()` を呼び出す

モデルのアーキテクチャと最終的なパラメーターをディスクに保存するなら、このタイミングが適しています。幅広い互換性を確保するため、モデルを Open Neural Network eXchange (ONNX) 形式で export します。そのファイル名を run.save() に渡すと、モデルのパラメーターが W&B のサーバーに保存されます。これで、どの .h5 や .pb がどのトレーニング run のものか分からなくなる心配はありません。モデルの保存、バージョン管理、配布に使える wandb のより高度な機能については、Artifacts ツールを参照してください。

def test(model, test_loader):
    model.eval()

    with wandb.init(project="pytorch-demo") as run:
        # テスト用のサンプルでモデルを実行します
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

            print(f"Accuracy of the model on the {total} " +
                f"test images: {correct / total:%}")
            
            run.log({"test_accuracy": correct / total})

        # 互換性の高い ONNX 形式でモデルを保存します
        torch.onnx.export(model, images, "model.onnx")
        run.save("model.onnx")

トレーニングを実行し、wandb.ai でメトリクスをリアルタイムに確認する

パイプライン全体を定義し、W&B のコードを数行追加したので、すべてが追跡される実験を実行する準備が整いました。 W&B からはいくつかのリンクが表示されます。ドキュメント、Project page (project 内のすべての run を整理するページ)、Run page (この run の結果が保存されるページ) です。 Run page にアクセスし、次のタブを確認してください。

Charts: トレーニング全体を通じて、モデルの勾配、パラメーター値、損失がログされます。
System: ディスク I/O 使用率、CPU および GPU のメトリクスなどのシステムメトリクスが表示されます。
Logs: トレーニング中に標準出力へ出力された内容のコピーがすべて含まれます。
Files: トレーニングが完了したら、model.onnx をクリックすると、Netron model viewer でネットワークを表示できます。

with wandb.init() ブロックを抜けて run が終了すると、W&B はセルの出力に結果のサマリーも表示します。

# パイプラインを使ってモデルの構築、トレーニング、分析を行う
model = model_pipeline(config)

sweep でハイパーパラメーターをテストする

この例では、1 組のハイパーパラメーターのみを扱いました。しかし、ほとんどの ML ワークフローでは、複数のハイパーパラメーターを試しながら反復することが重要です。 W&B Sweeps を使用すると、ハイパーパラメーターのテストを自動化し、候補となるモデルや最適化戦略の空間を探索できます。これにより、前述の単一設定による run にとどまらず、規模を拡大できます。 W&B Sweeps を使用したハイパーパラメーター最適化のデモを紹介する Colab ノートブックをご覧ください。 W&B でハイパーパラメーター sweep を実行する手順は、次の 3 ステップです。

sweep を定義する: 探索するパラメーター、探索戦略、最適化するメトリクスなどを指定する辞書または YAML ファイルを作成します。
sweep を初期化する: sweep_id = wandb.sweep(sweep_config)
sweep エージェントを実行する: wandb.agent(sweep_id, function=train)

ハイパーパラメーター sweep を実行するのに必要な作業はこれだけです。

サンプルギャラリー

W&B で追跡・可視化された project のサンプルは、ギャラリーでご覧いただけます。

高度な設定

以下のオプションを使用すると、前述の基本的なワークフローを本番環境、オフライン環境、またはマネージド環境向けに拡張できます。

環境変数: APIキーを環境変数に設定しておくと、マネージドクラスター上でトレーニングを実行できます。
オフラインモード: dryrun モードを使用すると、オフラインでトレーニングを行い、結果を後で同期できます。
オンプレミス: 自社インフラストラクチャー内のプライベートクラウドやエアギャップ環境のサーバーに W&B をインストールします。
Sweeps: チューニング用の軽量なツールで、ハイパーパラメーター探索をすばやく設定できます。

​このノートブックの内容

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

​ステップ 0: W&B をインストールする

​ステップ 1: W&B をインポートしてログインする

​実験とパイプラインを定義する

​wandb.init() でメタデータとハイパーパラメーターをトラッキングする

​データの読み込みとモデルを定義する

​トレーニングロジックを定義する

​run.watch() で勾配をトラッキングし、それ以外はすべて run.log() でトラッキングする

​テストロジックを定義する

​オプション: run.save() を呼び出す

​トレーニングを実行し、wandb.ai でメトリクスをリアルタイムに確認する

​sweep でハイパーパラメーターをテストする

​サンプルギャラリー

​高度な設定