PyTorch Ignite 0.4.8 : Tutorials : Getting Started on Cloud TPU (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 04/04/2022 (0.4.8)

* 本ページは、Pytorch Ignite の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Tutorials : Getting Started with PyTorch-Ignite on Cloud TPUs

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

 

 

Tutorials : Getting Started on Cloud TPU

このノートブックは “Getting Started with PyTorch on Cloud TPUs” に基づいて以下を行なう方法を示します :

• Colab 上に PyTorch/XLA をインストールします、これは貴方に TPU で PyTorch を利用させることができます。
• PyTorch-Ignite で MNIST 上で基本的なモデルを訓練します。

PyTorch/XLA は、PyTorch を Cloud TPU に接続させて TPU コアをデバイスとして使用させることができるパッケージです。Colab は無料の Cloud TPU システム (遠隔 CPU ホスト + それぞれ 2 コアを持つ 4 個の TPU チップ) を提供し、PyTorch/XLA のインストールは数分しかかかりません。

 

Colab Cloud TPU の使用

• メインメニューで、”Runtime” をクリックして “Change runtime type” を選択します。”hardware accelerator” として “TPU” を設定します。
• 下のセルは Colab 上で TPU へのアクセスを持つか確認します。

import os
assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'

 

PyTorch/XLA のインストール

PyTorch, Torchvision, と PyTorch/XLA をインストールするためには次のセルを実行します。実行には数分かかります。

PyTorch/XLA パッケージは PyTorch を Cloud TPU に接続させます。(XLA が TPU コンパイラの名前なので、PyTorch/TPU ではなく PyTorch/XLA と命名されています。) 特に、PyTorch/XLA は TPU コアを PyTorch デバイスとして利用可能にします。これは PyTorch を TPU 上でテンソルを作成して操作させます。

VERSION = !curl -s https://api.github.com/repos/pytorch/xla/releases/latest | grep -Po '"tag_name": "v\K.*?(?=")'
VERSION = VERSION[0].rstrip('0').rstrip('.') # remove trailing zero
!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-{VERSION}-cp37-cp37m-linux_x86_64.whl

 

必要な依存性

torch と ignite は既にインストールされていることを仮定しています。pip を使用してそれをインストールできます :

!pip install pytorch-ignite

 

PyTorch-Ignite で MNIST 上で基本的なモデルを訓練する

PyTorch XLA API は PyTorch が Cloud TPU を使用するとき丁度 CPU や CUDA デバイスを使用するように簡単です。わずかな変更だけで PyTorch と Ignite でモデルを訓練することができます。このサンプルのコードを使用します : https://github.com/pytorch/ignite/blob/master/examples/mnist/mnist_with_tensorboard_on_tpu.py

 

ライブラリのインポート

import torch
print("PyTorch version:", torch.__version__)

# imports the torch_xla package
import torch_xla
import torch_xla.core.xla_model as xm
print("PyTorch xla version:", torch_xla.__version__)

# Import PyTorch, Torchvision and Tensorboard
from torch.utils.data import DataLoader
from torch import nn
import torch.nn.functional as F
from torch.optim import SGD
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize

from torch.utils.tensorboard import SummaryWriter

# Import PyTorch-Ignite
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Accuracy, Loss, RunningAverage
from ignite.contrib.handlers import ProgressBar

 

データ処理

# Dataloaders
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_batch_size = 64
val_batch_size = train_batch_size * 2

train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)

 

モデルの作成

# Setup a basic CNN
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()
device = xm.xla_device()
model = model.to(device)  # Move model before creating optimizer

 

Optimizer と trainers

optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)

# Create trainer and evaluator
trainer = create_supervised_trainer(
    model, 
    optimizer, 
    F.nll_loss, 
    device=device, 
    output_transform=lambda x, y, y_pred, loss: [loss.item(), ]
)

evaluator = create_supervised_evaluator(
    model, metrics={"accuracy": Accuracy(), "nll": Loss(F.nll_loss)}, device=device
)

 

ハンドラ

# Setup event's handlers
log_interval = 10

log_dir = "/tmp/tb_logs"

# writer
writer = SummaryWriter(log_dir=log_dir)

tracker = xm.RateTracker()

# Add RateTracker as an output of the training step
@trainer.on(Events.ITERATION_COMPLETED)
def add_rate_tracker(engine):
    tracker.add(len(engine.state.batch))
    engine.state.output.append(tracker.global_rate())

# Setup output values of the training step as EMA metrics
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, "batch_loss")
RunningAverage(output_transform=lambda x: x[1]).attach(trainer, "global_rate")

# Let's log the EMA metrics every `log_interval` iterations
@trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
def log_training_loss(engine):
    writer.add_scalar("training/batch_loss", engine.state.metrics["batch_loss"], engine.state.iteration)
    writer.add_scalar("training/global_rate", engine.state.metrics["global_rate"], engine.state.iteration)

# Setup a progress bar (tqdm) and display batch loss metric in the bar
pbar = ProgressBar()
pbar.attach(trainer, ["batch_loss", "global_rate"])

# Let's compute training metrics: average accuracy and loss
@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"]
    pbar.log_message(
        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)

# Let's compute training metrics: average accuracy and loss
@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)

# Display in Firefox may not work properly. Use Chrome.
%load_ext tensorboard

%tensorboard --logdir="/tmp/tb_logs"

# kick everything off
!rm -rf /tmp/tb_logs/*

trainer.run(train_loader, max_epochs=10)

以上