PyTorch

Try in Colab

機械学習の 実験管理 、データセットの バージョン管理 、および プロジェクト の共同作業に W&B を活用しましょう。

​

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

このチュートリアルでは、W&BをPyTorchコードに統合し、パイプラインに 実験管理 機能を追加する方法を紹介します。

# ライブラリのインポート
import wandb

# 新しい実験（experiment）を開始
with wandb.init(project="new-sota-model") as run:
 
    # configを使ってハイパーパラメーターの辞書をキャプチャ
    run.config = {"learning_rate": 0.001, "epochs": 100, "batch_size": 128}

    # モデルとデータのセットアップ
    model, dataloader = get_model(), get_data()

    # オプション：勾配（gradients）をトラック
    run.watch(model)

    for batch in dataloader:
    metrics = model.training_step()
    # モデルのパフォーマンスを可視化するために、トレーニングループ内でメトリクスをログ記録
    run.log(metrics)

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

ビデオチュートリアル もあわせてご覧ください。

注意: 既存のパイプラインにW&Bを統合するために必要なのは、 Step で始まるセクションだけです。残りの部分はデータのロードとモデルの定義を行っています。

​

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

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

# 決定論的な振る舞いを保証
torch.backends.cudnn.deterministic = True
random.seed(hash("setting random seeds") % 2**32 - 1)
np.random.seed(hash("improves reproducibility") % 2**32 - 1)
torch.manual_seed(hash("by removing stochasticity") % 2**32 - 1)
torch.cuda.manual_seed_all(hash("so runs are repeatable") % 2**32 - 1)

# デバイス設定
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のインストール

まずはライブラリを入手する必要があります。 wandb は pip を使って簡単にインストールできます。

!pip install wandb onnx -Uq

​

ステップ 1: W&Bのインポートとログイン

Webサービスにデータをログ記録するには、ログインする必要があります。

W&Bを初めて使用する場合は、表示されるリンクから無料アカウントを登録してください。

import wandb

wandb.login()

​

実験とパイプラインの定義

​

wandb.init でメタデータとハイパーパラメーターをトラックする

プログラム的にまず行うのは、実験の定義です。 ハイパーパラメーターは何か？ この run に関連付けられている メタデータ は何か？

これらの情報を config 辞書（または同様の オブジェクト ）に保存し、必要に応じてアクセスするのが一般的な ワークフロー です。

この例では、一部の ハイパーパラメーター のみを変更可能にし、残りはハードコードしています。しかし、 モデル のあらゆる部分を config に含めることができます。

また、MNIST データセット と 畳み込み アーキテクチャー を使用しているという メタデータ も含めています。後で同じ プロジェクト 内で CIFAR の全結合 アーキテクチャー を扱う場合などに、これによって Runs を区別しやすくなります。

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

次に、モデルトレーニングで一般的な全体のパイプラインを定義しましょう。

1. モデル、および関連する データ と オプティマイザー を make（作成）し、
2. それに応じてモデルを train（トレーニング）し、最後に
3. 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 に渡すと、それらの情報が即座にログ記録されます。これにより、実験で使用した ハイパーパラメーター の値を常に把握できます。

選択してログに記録した値が常に モデル で使用されることを確実にするために、オブジェクトのコピーである run.config を使用することをお勧めします。以下の make の定義で例を確認してください。

補足: 私たちはコードを別 プロセス で実行するように配慮しています。そのため、こちら側で問題が発生しても（巨大な海獣がデータセンターを襲ったとしても）、お客様のコードがクラッシュすることはありません。問題が解決した後（クラーケンが深海に戻った後など）に、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 アーキテクチャー を使用します。

自由にモデルを調整して実験を試してみてください。すべての 結果 は wandb.ai にログ記録されます。

# 標準的な畳み込みニューラルネットワーク

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（トレーニング）の方法を指定しましょう。

ここでは watch と log という2つの wandb 関数が登場します。

​

run.watch() で勾配を、それ以外を run.log() でトラックする

run.watch は、トレーニングの log_freq ステップごとに、モデルの 勾配（gradients）と パラメータ をログ記録します。

トレーニングを開始する前にこれを呼び出すだけです。

残りのトレーニングコードは同じです。 エポック と バッチ をイテレートし、 forward pass と backward pass を実行して オプティマイザー を適用します。

def train(model, loader, criterion, optimizer, config):
    # wandbにモデルの動作（勾配、重みなど）を監視（watch）するよう指示
    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)
    
    # Forward pass ➡
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # Backward pass ⬅
    optimizer.zero_grad()
    loss.backward()

    # オプティマイザーによる更新
    optimizer.step()

    return loss

唯一の違いはログ記録のコードにあります。以前は ターミナル に出力して メトリクス を報告していたかもしれませんが、代わりに同じ情報を run.log() に渡します。

run.log() は文字列を キー とする 辞書 を期待します。これらの文字列は、ログ記録される値である オブジェクト を識別します。オプションで、現在どのトレーニング step にいるかをログ記録することもできます。

補足: 私はモデルが処理したサンプル数を使用するのが好きです。これにより、バッチサイズ が異なっても比較が容易になります。もちろん、生のステップ数やバッチ数を使用することもできます。長時間のトレーニングの場合は、 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 がどのトレーニング Runs に対応するか分からなくなることはもうありません。

モデルの保存、 バージョン管理 、配布のためのより高度な 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コードを挿入したので、完全にトラックされた実験を実行する準備が整いました。

いくつかのリンクが表示されます： ドキュメント、 プロジェクト 内のすべての Runs を整理する Project ページ、そしてこの run の 結果 が保存される Run ページです。

Run ページに移動して、以下のタブを確認してください：

1. Charts: トレーニング中にモデルの 勾配 、 パラメータ 値、損失がログ記録されます。
2. System: ディスク I/O 使用率、 CPU および GPU メトリクス（温度の上昇を確認してください）など、さまざまなシステムメトリクスが含まれます。
3. Logs: トレーニング中に標準出力に送られたすべての内容のコピーが含まれます。
4. Files: トレーニング完了後、 model.onnx をクリックして Netron モデルビューアー でネットワークを表示できます。

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

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

​

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

この例では1セットの ハイパーパラメーター のみを確認しました。しかし、多くの機械学習 ワークフロー において重要なのは、複数の ハイパーパラメーター を反復試行することです。

W&B Sweeps を使用すると、 ハイパーパラメーター テストを自動化し、可能な モデル や最適化戦略の空間を探索できます。

W&B Sweeps を使用したハイパーパラメーター最適化の Colabノートブック をご覧ください。

W&Bでの ハイパーパラメーター探索（sweep） の実行は非常に簡単です。シンプルな3つのステップで行えます。

1. sweep を定義する: 探索する パラメータ 、探索戦略、最適化 メトリクス などを指定する 辞書 または YAML ファイル を作成します。

2. sweep を初期化する: sweep_id = wandb.sweep(sweep_config)

3. sweep agent を実行する: wandb.agent(sweep_id, function=train)

ハイパーパラメーター探索の実行に必要なのはこれだけです。

​

ギャラリー

W&Bでトラック・可視化された プロジェクト の例を ギャラリー → でご覧いただけます。

​

高度なセットアップ

1. 環境変数: 環境変数に APIキー を設定して、管理された クラスター でトレーニングを実行できます。
2. オフラインモード: dryrun モードを使用してオフラインでトレーニングし、後で 結果 を同期できます。
3. オンプレミス（On-prem）: 自社 インフラストラクチャー 内の プライベートクラウド やエアギャップのある サーバー にW&Bをインストールできます。個人研究者からエンタープライズチームまで利用可能なローカルインストールを提供しています。
4. Sweeps: チューニングのための軽量な ツール を使用して、 ハイパーパラメーター 探索を迅速にセットアップできます。