PyTorch Lightning 1.1 : 概要 (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 02/01/2021 (1.1.x)

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

• README

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

PyTorch Lightning 1.1 : 概要

高パフォーマンス AI 研究のための軽量 PyTorch ラッパーです。モデルをスケールし、ボイラープレートではありません。

ニュース

• Dec 2020 – Facebook が研究とプロダクション・チームに渡り深層学習を標準化するためにどのように Lightning を利用するかについて読んでください

PyTorch Lightning は単なる体系化された PyTorch です

Lightning はエンジニアリングからサイエンスを分離するために PyTorch コードのもつれを解きます。

Lightning 哲学

Lightning はこれらの原則を念頭において設計されています :

原則 1 : 最大限の柔軟性を可能にする。原理 2 : 不要なボイラープレートを取り去りますが、必要なときにはアクセス可能にします。原理 3 : システムは自己充足的であるべきです (i.e. optimizer、計算コード等)。原理 4 : 深層学習コードは 4 つの異なるカテゴリーに体系化されるべきです。

• 研究コード (LightningModule)。
• エンジニアリング・コード ( 除去し、そして Trainer により処理されます)。
• 本質的でない研究コード (ロギング等… これは Callback に行きます)。
• データ (PyTorch Dataloader を使用するかそれらを LightningDataModule に体系化します)。

ひとたびこれを行えば、貴方のコードを変更することなくマルチ GPU, TPU, CPU 上、そして 16-bit precision 内でさえ訓練できます。

2 ステップガイド で始めてください。

推論

Lightning はまた AI 研究者とプロダクション・チームが BERT と自己教師あり学習のようなものをスケールアップする必要がある高速な推論のためにも設計されています。Lightning はそれらのケースのために ONNX や TorchScript に自動的にエクスポートできます。

How To Use

ステップ 0: インストール

PyPI からの単純なインストール :

pip install pytorch-lightning

完全なパッケージ経験を得るには、pytorch-lightning[‘extra’] あるいは CPU ユーザのためには pytorch-lightning[‘cpu-extra’] で総てのオプションの依存性をインストールすることもできます。

From Conda

conda install pytorch-lightning -c conda-forge

ステップ 1 : これらのインポートを追加する

import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import pytorch_lightning as pl

ステップ 2 : LightningModule (nn.Module サブクラス) を定義する

LightningModule は完全なシステム (ie: a GAN, オートエンコーダ, BERT or 単純な画像分類器) を定義します。

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defined the train loop. It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

Note: training_step は訓練ループを定義します。forward は推論/予測の間に LightningModule がどのように動作するかを定義します。

ステップ 3 : 訓練！

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train, val = random_split(dataset, [55000, 5000])

autoencoder = LitAutoEncoder()
trainer = pl.Trainer()
trainer.fit(autoencoder, DataLoader(train), DataLoader(val))

そしてコードの単一行を変更することなしに、GPU/TPU 上で実行できるでしょう。

# 8 GPUs
trainer = Trainer(max_epochs=1, gpus=8)

# 256 GPUs
trainer = Trainer(max_epochs=1, gpus=8, num_nodes=32)

# TPUs
trainer = Trainer(tpu_cores=8)

そして onnx や torchscript を通してプロダクションのためにエクスポートさえできるでしょう。

# torchscript
autoencoder = LitAutoEncoder()
torch.jit.save(autoencoder.to_torchscript(), "model.pt")

# onnx
with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
    autoencoder = LitAutoEncoder()
    input_sample = torch.randn((1, 64))
    autoencoder.to_onnx(tmpfile.name, input_sample, export_params=True)
    os.path.isfile(tmpfile.name)

上級ユーザのためには、依然として複雑な訓練ループを所有できます。

class LitAutoEncoder(pl.LightningModule):
    def training_step(self, batch, batch_idx, optimizer_idx):
        # access your optimizers with use_pl_optimizer=False. Default is True
        (opt_a, opt_b) = self.optimizers(use_pl_optimizer=True)

        loss_a = ...
        self.manual_backward(loss_a, opt_a)
        opt_a.step()
        opt_a.zero_grad()

        loss_b = ...
        self.manual_backward(loss_b, opt_b, retain_graph=True)
        self.manual_backward(loss_b, opt_b)
        opt_b.step()
        opt_b.zero_grad()

主要な特徴

• モデルを変更することなく任意のハードウェア (CPU, GPU, TPU) 上で実行するために貴方のモデルをスケールできます。
• 研究コードをエンジニアリングから切り離すことによりコードの可読性を高めます。
• より容易に再現する
• 訓練ループとトリッキーなエンジニアリングを自動化することによりエラーを少ない傾向にします。
• 総ての柔軟性を保持しながら (LightningModule は依然として PyTorch モジュールです)、多くのボイラープレートは取り除きます。
• Lightning はポピュラーなロギング/可視化フレームワーク (Tensorboard, MLFlow, Neptune.ai, Comet.ml, Wandb) との out-of-the-box な統合を持ちます。
• 総ての新しい PR で厳密にテストしました。PyTorch と Python サポート・バージョン、総ての OS、マルチ GPU そして TPU とさえも総ての組合せをテストしています。
• 最小限の実行スピードのオーバーヘッド (純粋な PyTorch に比べてエポック毎に約 300 ms)。

Lightning は DL/ML 研究の 40+ パーツを自動化する

• GPU 訓練
• 分散 GPU (クラスタ) 訓練
• TPU 訓練
• EarlyStopping
• ロギング / 可視化
• チェックポインティング
• 実験管理
• 完全なリストはここです

以上