PyTorch Ignite 0.4.2 : Examples : MNIST with TensorBoard (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 02/09/2021 (0.4.2)

* 本ページは、PyTorch Ignite ドキュメントの以下のサンプルを適宜書き換えて補足説明したものです：

• examples/mnist/mnist_with_tensorboard.py

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

PyTorch Ignite 0.4.2 : Examples : MNIST with TensorBoard

TensorBoard を利用して訓練と検証を監視する MNIST サンプルです。

from argparse import ArgumentParser

import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, Normalize, ToTensor

from ignite.engine import Events, create_supervised_evaluator, create_supervised_trainer
from ignite.metrics import Accuracy, Loss

from torch.utils.tensorboard import SummaryWriter

train_batch_size = 64
val_batch_size = 1000
epochs = 10
lr = 0.01
momentum = 0.5
log_interval = 10
log_dir = "tensorboard_logs"

訓練と検証データセットを torch.utils.data.DataLoader として定義します :

def get_data_loaders(train_batch_size, val_batch_size):
    data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])

    train_loader = DataLoader(
        MNIST(download=True, root=".", transform=data_transform, train=True), batch_size=train_batch_size, shuffle=True
    )

    val_loader = DataLoader(
        MNIST(download=False, root=".", transform=data_transform, train=False), batch_size=val_batch_size, shuffle=False
    )
    return train_loader, val_loader

train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)

device = "cpu"
if torch.cuda.is_available():
  device = "cuda"

モデルを定義します :

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=-1)

model = Net()
model.to(device)  # Move model before creating optimizer

optimizer と損失関数を定義します :

optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
criterion = nn.NLLLoss()

次に trainer と evaluator エンジンを定義します。このサンプルではヘルパー・メソッド create_supervised_trainer() と create_supervised_evaluator() を使用しています :

trainer = create_supervised_trainer(model, optimizer, F.nll_loss, device=device)

evaluator を作成するヘルパー関数 create_supervised_evaluator は引数 metrics を受け取ることに注意してください :

val_metrics = {"accuracy": Accuracy(), "nll": Loss(criterion)}
evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)

ここで 2 つのメトリクスを定義しています : 検証データ・セット上で計算するための精度と損失です。メトリクスのより多くの情報は ignite.metrics で見つかります。

オブジェクト trainer と evaluator は Engine のインスタンスです – Ignite の主要コンポーネントです。Engine は訓練/検証ループに渡る抽象です。
※ 一般に、Engine クラスとカスタム訓練/検証ステップ・ロジックを直接使用して trainer と evaluator を定義することができます。

 
TensorBoard を利用するために SummaryWriter のインスタンスを生成します :

writer = SummaryWriter(log_dir=log_dir)

コードスニペットの最も興味深いパートはイベント・ハンドラの追加です。Engine は実行の間にトリガーされる様々なイベント上にハンドラを追加することを可能にします。イベントがトリガーされたとき、装着されたハンドラ (関数) が実行されます。そして、ロギング目的で総ての log_interval -th 反復の終わりに実行される関数を追加しました :

@trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
def log_training_loss(engine):
    print(
        f"Epoch[{engine.state.epoch}] Iteration[{engine.state.iteration}/{len(train_loader)}] "
        f"Loss: {engine.state.output:.2f}"
    )
    writer.add_scalar("training/loss", engine.state.output, engine.state.iteration)

エポックが終了するとき訓練と検証メトリクス (*1) を計算することを望みます。その目的で train_loader と val_loader 上で前に定義した evaluator を実行できます。そのため epoch complete イベント上で trainer に 2 つの追加のハンドラを装着できます :

@trainer.on(Events.EPOCH_COMPLETED)
def log_training_results(engine):
    evaluator.run(train_loader)
    metrics = evaluator.state.metrics
    avg_accuracy = metrics["accuracy"]
    avg_nll = metrics["nll"]
    print(
        f"Training Results - Epoch: {engine.state.epoch} Avg accuracy: {avg_accuracy:.2f} Avg loss: {avg_nll:.2f}"
    )
    writer.add_scalar("training/avg_loss", avg_nll, engine.state.epoch)
    writer.add_scalar("training/avg_accuracy", avg_accuracy, engine.state.epoch)

@trainer.on(Events.EPOCH_COMPLETED)
def log_validation_results(engine):
    evaluator.run(val_loader)
    metrics = evaluator.state.metrics
    avg_accuracy = metrics["accuracy"]
    avg_nll = metrics["nll"]
    print(
        f"Validation Results - Epoch: {engine.state.epoch} Avg accuracy: {avg_accuracy:.2f} Avg loss: {avg_nll:.2f}"
    )
    writer.add_scalar("valdation/avg_loss", avg_nll, engine.state.epoch)
    writer.add_scalar("valdation/avg_accuracy", avg_accuracy, engine.state.epoch)

最後に、訓練データセット上でエンジンをスタートさせて 100 エポックの間それを実行します :

# kick everything off
trainer.run(train_loader, max_epochs=epochs)

writer.close()

以上