PyTorch

Colab で試す

機械学習の実験トラッキング、データセットのバージョン管理、プロジェクトでの共同作業に W&B を利用します。

​

このノートブックで扱う内容

PyTorch のコードと W&B をインテグレーションし、パイプラインに実験トラッキングを追加する方法を説明します。

# ライブラリをインポート
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")

ビデオチュートリアルを見ながら進めてください。

注記: 「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

# 決定論的な動作を保証する
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

​

Step 1: W&B をインポートしてログインする

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

初めて W&B を使用する場合は、 表示されるリンクから無料アカウントを作成する必要があります。

import wandb

wandb.login()

​

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

​

wandb.init を使ってメタデータとハイパーパラメータを追跡する

プログラム上で最初に行うのは、実験を定義することです。 どのハイパーパラメータを使うのか？この run にはどんなメタデータが紐づいているのか？

この情報を config 辞書（または類似のオブジェクト）に保存し、 必要に応じてそこから参照する、というワークフローはとても一般的です。

この例では、一部のハイパーパラメータだけを変化させ、 残りはコードに手書きで埋め込んでいます。 ただし、モデルの任意の部分を config に含めることができます。

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

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 のコンテキスト内で実行されるということだけです。 この関数を呼び出すと、あなたのコードと当社サーバーとの 通信経路が確立されます。

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 するかを指定します。

ここでは 2 つの wandb 関数である watch と log を使用します。

​

run.watch() で勾配をトラッキングし、その他は run.log() で記録する

run.watch は、学習の log_freq ステップごとに、 モデルの勾配とパラメータをログに記録します。

学習を開始する前にこれを呼び出すだけでかまいません。

それ以外の学習コードはそのままです： エポックとバッチを反復しながら、 forward / backward パスを実行し、 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() はキーとして文字列を持つ辞書を受け取ります。 これらの文字列でログされるオブジェクトを識別し、そのオブジェクトが値になります。 またオプションとして、今どの step の学習中かをログすることもできます。

補足: 私は、モデルがこれまでに見たサンプル数を使うのが好きです。 こうするとバッチサイズが異なる場合でも比較しやすくなるからです。 ただし、生のステップ数やバッチ数を使ってもかまいません。学習 run が長くなる場合には、エポック ごとにログを取るのも理にかなっています。

def train_log(loss, example_ct, epoch):
    with wandb.init(project="pytorch-demo") as run:
        # lossとエポック数をログに記録する
        # ここでメトリクスを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) format でモデルを export します。

そのファイル名を run.save() に渡すことで、モデルのパラメータが W&B のサーバーに保存されるようになります。どの .h5 や .pb がどの学習 run に対応しているのか分からなくなることは、もうありません。

モデルの保存、バージョン管理、配布のための、より高度な wandb の機能については、Artifacts tools を参照してください。

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 コードを追加したので、 これで完全にトラッキングされる実験を実行する準備ができました。

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

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

1. Charts: 学習全体を通して、モデルの勾配、パラメータ値、損失がログされます
2. System: Disk I/O 使用率、CPU と GPU のメトリクス（温度の上昇に注意）など、さまざまなシステムメトリクスが含まれます
3. Logs: 学習中に標準出力へ出力された内容のコピーが含まれます
4. Files: 学習が完了すると、model.onnx をクリックして Netron model viewer でネットワークを表示できます。

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

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

​

Sweeps を使ってハイパーパラメータをテストする

この例では、1 つのハイパーパラメータセットだけを扱いました。 しかし、ほとんどの ML ワークフローでは、 複数のハイパーパラメータを繰り返し試すことが重要です。

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

W&B Sweeps を使ったハイパーパラメータ最適化をデモする Colab ノートブックを参照してください。

W&B を使ってハイパーパラメータスイープを実行するのは非常に簡単です。たった 3 つのステップです。

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

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

3. スイープエージェントを実行する: wandb.agent(sweep_id, function=train)

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

​

サンプルギャラリー

W&B でトラッキングおよび可視化されているプロジェクトの例は、Gallery →でご覧いただけます。

​

高度なセットアップ

1. 環境変数: 管理されたクラスタ上で学習を実行できるように、APIキーを環境変数に設定します。
2. オフラインモード: 結果を後から同期できるように、dryrun モードを使ってオフラインで学習します。
3. オンプレミス: 自社インフラストラクチャ内のプライベートクラウドまたはエアギャップされたサーバーに W&B をインストールします。当社では、研究機関からエンタープライズの Teams まで、あらゆるお客様向けにローカルインストールを提供しています。
4. Sweeps: 軽量なチューニングツールを使って、ハイパーパラメータ探索をすばやくセットアップします。