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

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

• examples/mnist/mnist.py

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

PyTorch Ignite 0.4.2 : Examples : MNIST

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 tqdm import tqdm

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

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

訓練と検証データセットを 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)

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./MNIST/raw/train-images-idx3-ubyte.gz
72%
7127040/9912422 [00:00<00:15, 184850.16it/s]
Extracting ./MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/raw/train-labels-idx1-ubyte.gz
32768/? [00:00<00:00, 110253.88it/s]
Extracting ./MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/raw/t10k-images-idx3-ubyte.gz
1654784/? [00:00<00:00, 7488051.18it/s]
Extracting ./MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/raw/t10k-labels-idx1-ubyte.gz


0/? [00:00<?, ?it/s]
Extracting ./MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/raw
Processing...
Done!
/usr/local/lib/python3.6/dist-packages/torchvision/datasets/mnist.py:480: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:141.)
  return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)

%ls -l MNIST/raw

total 65008
-rw-r--r-- 1 root root  7840016 Feb  7 08:14 t10k-images-idx3-ubyte
-rw-r--r-- 1 root root  1648877 Feb  7 08:14 t10k-images-idx3-ubyte.gz
-rw-r--r-- 1 root root    10008 Feb  7 08:14 t10k-labels-idx1-ubyte
-rw-r--r-- 1 root root     4542 Feb  7 08:14 t10k-labels-idx1-ubyte.gz
-rw-r--r-- 1 root root 47040016 Feb  7 08:14 train-images-idx3-ubyte
-rw-r--r-- 1 root root  9912422 Feb  7 08:14 train-images-idx3-ubyte.gz
-rw-r--r-- 1 root root    60008 Feb  7 08:14 train-labels-idx1-ubyte
-rw-r--r-- 1 root root    28881 Feb  7 08:14 train-labels-idx1-ubyte.gz

%ls -l MNIST/processed

total 54144
-rw-r--r-- 1 root root  7921079 Feb  7 08:14 test.pt
-rw-r--r-- 1 root root 47521079 Feb  7 08:14 training.pt

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, criterion, device=device)
trainer.logger = setup_logger("trainer")

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

val_metrics = {"accuracy": Accuracy(), "nll": Loss(criterion)}

evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)
evaluator.logger = setup_logger("evaluator")

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

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

def train_step(engine, batch):
    model.train()
    optimizer.zero_grad()
    x, y = batch[0].to(device), batch[1].to(device)
    y_pred = model(x)
    loss = criterion(y_pred, y)
    loss.backward()
    optimizer.step()
    return loss.item()

trainer = Engine(train_step)

def validation_step(engine, batch):
    model.eval()
    with torch.no_grad():
        x, y = batch[0].to(device), batch[1].to(device)
        y_pred = model(x)
        return y_pred, y

evaluator = Engine(validation_step)

pbar = tqdm(initial=0, leave=False, total=len(train_loader), desc=f"ITERATION - loss: {0:.2f}")

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

@trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
def log_training_loss(engine):
  pbar.desc = f"ITERATION - loss: {engine.state.output:.2f}"
  pbar.update(log_interval)

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

@trainer.on(Events.EPOCH_COMPLETED)
def log_training_results(engine):
    pbar.refresh()
    evaluator.run(train_loader)
    metrics = evaluator.state.metrics
    avg_accuracy = metrics["accuracy"]
    avg_nll = metrics["nll"]
    tqdm.write(
        f"Training Results - Epoch: {engine.state.epoch} Avg accuracy: {avg_accuracy:.2f} Avg loss: {avg_nll:.2f}"
    )

@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"]
    tqdm.write(
        f"Validation Results - Epoch: {engine.state.epoch} Avg accuracy: {avg_accuracy:.2f} Avg loss: {avg_nll:.2f}"
    )

    pbar.n = pbar.last_print_n = 0

@trainer.on(Events.EPOCH_COMPLETED | Events.COMPLETED)
def log_time(engine):
    tqdm.write(f"{trainer.last_event_name.name} took { trainer.state.times[trainer.last_event_name.name]} seconds")

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

trainer.run(train_loader, max_epochs=epochs)

pbar.close()

以上