PyTorch 三国志（Ignite・Catalyst・Lightning）

(2021/06/19) 2021年の PyTorch Lightning については PyTorch Lightning 2021 (for MLコンペ) - Qiita にまとめました。

0. 導入

深層学習フレームワークはいずれも開発が非常に速く盛り上がっている分野だと思います。
TensorFlow や jax 等もある中、つい先日 PFN のニュースもあり、PyTorch もより盤石となりそうです。おそらくこれからも PyTorch ユーザーは増えると思われます(Chainer にもあった公式 Trainer が PyTorch 内に実装されるとこの記事の存在が危ぶまれるので、そこには触れないこととします)。

しかし PyTorch は自由度が高い一方、学習周りのコード(各 epoch のループ周りとか)は個々人に委ねられており、非常に個性豊かなコードとなりがちです。

これらのコードを自分で書くことは非常に学びが多く、PyTorch を始める場合には必ず通るべきだと私は思います。しかしあまりに個性が強すぎると、他の人との共有やコンペ間での使い回し等のシーンで辛いときがあります(ex. Winner Solutions でよく見かけるオレオレ Trainer)。

PyTorch の場合、学習周りのコードを簡略化するためのフレームワークは自身の中にはない(以前 Trainer があったが廃止された)のですが、Ecosystem | PyTorch の中では以下の PyTorch 用フレームワークが紹介されています。

• catalyst
• fastai
• ignite
• pytorch-lightning

多いですね。
全部試して自分に合ったものを見つけろというのは正論です。しかしそれらは楽ではないので本記事では各フレームワークの紹介と簡単な比較をしてみて、皆さんが触ってみる何かしらの目安になればと思います。

なお fastai については頭一つ抜けて抽象度が高い(コードが短くなりやすい)のですが、自身で細かい操作を加えるための学習コストが高く感じたため、本記事内の比較では予め省いております。

そのため本記事では Catalyst・Ignite・Lightning の3つに絞り、かつ Kaggle のコンペに参加することを想定して比較を行っていきます。

ちなみにこれらの 3つのフレームワークについては予めある程度動作することは確認しました。本記事を読んでもう少し踏み込みたくなった方はご参照いただければと思います。

• 本記事内のコード
  • https://github.com/yukkyo/Compare-PyTorch-Catalyst-Ignite-Lightning
• Catalyst: yukkyo/PyTorch-FilterResponseNormalizationLayer
• Ignite: PyTorch-Ignite で学習用コードをスマートにする - ふぁむたろうのブログ
• Lightning: yukkyo/Kaggle-Understanding-Clouds-69th-solution

先に述べますが、いずれもコンペに参加できるだけのポテンシャルはあります。

1. この記事の対象(とか対象外)

• Kaggle コンペに興味ある
• 画像系コンペに興味ある
  • 特に Classification・Segmentation・Detection あたりに興味ある
• PyTorch 触ったことある
  • 触ったことない人はこの記事読んでる場合じゃないです
  • 以下のページや本とかで始めましょう
    • Welcome to PyTorch Tutorials — PyTorch Tutorials 1.3.1 documentation
    • つくりながら学ぶ！PyTorchによる発展ディープラーニング | 小川 雄太郎 | 工学 | Kindleストア | Amazon

2. 各フレームワーク(Catalyst・Ignite・Lightning)比較

2019年12月13日時点の pip 上での最新版を使いました。
Python のバージョンは 3.7.5 です。また NVIDIA/apexもインストール済を想定しています。
このコードを動かす分には apex は不要です。

torch==1.3.1
torchvision==0.4.2
catalyst==19.12
pytorch-ignite==0.2.1
pytorch-lightning==0.5.3.2

2.1 Star 数遷移(2019年12月10日時点)

Catalyst・Ignite が順調に伸びている一方、Lightning は今年4月からすごい勢いで伸びてきました。
一方 Lightning はまだ世に出て一年も経っていないので開発中の機能も多く、まだ unstable(バージョン上げたときに後方互換性がない等)であることには注意です。

また最近の Kaggle Notebook 上では Catalyst をよく見かけるため、Catalyst が Ignite を追い抜かすこともありえそうです。

2.2 書き方

ここでは素の PyTorch 学習用コードに対し各フレームワークを適用したらどうなるのか確認します。

2.2.1 共通部分

今回は cifer10 dataset に対して Resnet18 で学習してみようと思います。
下記のコードのように、モデルや Dataloader の定義は予め関数にしておきます。

共通部分のコード(長いのでたたみました)

share_funcs.py

import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
from torchvision import datasets, models, transforms

def get_criterion():
    """Loss をよしなに返してくれる関数"""
    return nn.CrossEntropyLoss()

def get_loaders(batch_size: int = 16, num_workers: int = 4):
    """各 Dataloader をよしなに返してくれえる関数"""
    transform = transforms.Compose([transforms.ToTensor()])

    # Dataset
    args_dataset = dict(root='./data', download=True, transform=transform)
    trainset = datasets.CIFAR10(train=True, **args_dataset)
    testset = datasets.CIFAR10(train=False, **args_dataset)

    # Data Loader
    args_loader = dict(batch_size=batch_size, num_workers=num_workers)

    train_loader = DataLoader(trainset, shuffle=True, **args_loader)
    val_loader = DataLoader(testset, shuffle=False, **args_loader)
    return train_loader, val_loader

def get_model(num_class: int = 10):
    """モデルをよしなに返してくれる関数"""
    model = models.resnet18(pretrained=True)
    num_features = model.fc.in_features
    model.fc = nn.Linear(num_features, num_class)
    return model

def get_optimizer(model: torch.nn.Module, init_lr: float = 1e-3, epoch: int = 10):
    optimizer = optim.SGD(model.parameters(), lr=init_lr, momentum=0.9)
    lr_scheduler = optim.lr_scheduler.MultiStepLR(
        optimizer,
        milestones=[int(epoch*0.8), int(epoch*0.9)],
        gamma=0.1
    )
    return optimizer, lr_scheduler

2.2.1 ベースコード(素の学習用コード)

あまり深く考えずに愚直に書くと下記のようになると思います。
.to(device) や loss.backward()、optimizer.step() は書かなきゃいけないので、どうしても長くなりがちです。
また with torch.no_grad() は torch.set_grad_enabled(bool) を使うことで Train と Eval 時の両方に対応させることは可能なのですが、Train と Eval 時は違う処理が多く(ex. optimizer.step() や metrics 等)、両方対応させるような関数を作るとかえって見通しが悪くなりがちです。

ベースコード(長いのでたたみました)

def train(model, data_loader, criterion, optimizer, device, grad_acc=1):
    model.train()

    # zero the parameter gradients
    optimizer.zero_grad()

    total_loss = 0.
    for i, (inputs, labels) in tqdm(enumerate(data_loader), total=len(data_loader)):
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)

        loss = criterion(outputs, labels)
        loss.backward()

        # Gradient accumulation
        if (i % grad_acc) == 0:
            optimizer.step()
            optimizer.zero_grad()

        total_loss += loss.item()

    total_loss /= len(data_loader)
    metrics = {'train_loss': total_loss}
    return metrics


def eval(model, data_loader, criterion, device):
    model.eval()
    num_correct = 0.

    with torch.no_grad():
        total_loss = 0.
        for inputs, labels in tqdm(data_loader, total=len(data_loader)):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            loss = criterion(outputs, labels)

            total_loss += loss.item()
            num_correct += torch.sum(preds == labels.data)

        total_loss /= len(data_loader)
        num_correct /= len(data_loader.dataset)
        metrics = {'valid_loss': total_loss, 'val_acc': num_correct}
    return metrics


def main():
    epochs = 10

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = get_model()
    train_loader, val_loader = get_loaders()
    optimizer, lr_scheduler = get_optimizer(model=model)
    criterion = get_criterion()

    # Model を multi-gpu したり、FP16 対応したりする
    model = model.to(device)

    print('Train start !')
    for epoch in range(epochs):
        print(f'epoch {epoch} start !')
        metrics_train = train(model, train_loader, criterion, optimizer, device)
        metrics_eval = eval(model, val_loader, criterion, device)

        lr_scheduler.step()

        # Logger 周りの処理
        # print するためのごちゃごちゃした処理
        print(f'epoch: {epoch} ', metrics_train, metrics_eval)

        # tqdm 使ってたらさらにごちゃごちゃする処理をここに書く
        # Model を保存するための処理
        # Multi-GPU の場合さらに注意して書く

2.2.2 Catalyst

Catalyst の場合、ライブラリ内の SupervisedRunner に必要なものを渡せば終わりです。すごいスマートですね！
また Accuracy や Dice 等のメジャーな metrics であれば Catalyst 内にあるため、それらを使えば自分で書くことはほとんどありません(独自 metrics の導入も比較的楽そうでした)。
大抵デフォルトのままで困らなそうですが、自分で細かい処理を加えたい場合若干調べる必要がありそうです。

import catalyst
from catalyst.dl import SupervisedRunner
from catalyst.dl.callbacks import AccuracyCallback
from share_funcs import get_model, get_loaders, get_criterion, get_optimizer

def main():
    epochs = 5
    num_class = 10
    output_path = './output/catalyst'

    model = get_model()
    train_loader, val_loader = get_loaders()
    loaders = {"train": train_loader, "valid": val_loader}

    optimizer, lr_scheduler = get_optimizer(model=model)
    criterion = get_criterion()

    runner = SupervisedRunner(device=catalyst.utils.get_device())
    runner.train(
        model=model,
        criterion=criterion,
        optimizer=optimizer,
        scheduler=lr_scheduler,
        loaders=loaders,
        logdir=output_path,
        callbacks=[AccuracyCallback(num_classes=num_class, accuracy_args=[1])],
        num_epochs=epochs,
        main_metric="accuracy01",
        minimize_metric=False,
        fp16=None,
        verbose=True
    )

2.2.3 Ignite

Ignite は Catalyst や後述する Lightning とは少し毛色が違います。
下記のように @trainer.on(Events.EPOCH_COMPLETED) 等で自分が挟みたい処理を各タイミングに対して差し込んでいくようなイメージです。
また Ignite も公式で Accuracy 等は用意されているのでメジャーな評価指標であれば自分で定義せずに済みそうです。

一方使いこなすのに慣れが必要そうなのと、イベント挟み方の自由度が高い(trainder.append のような足し方もできる)ので、一歩間違えると全体の見通しが悪くなる可能性もあります。

import torch
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Accuracy, Loss, RunningAverage
from ignite.contrib.handlers import ProgressBar
from share_funcs import get_model, get_loaders, get_criterion, get_optimizer

def run(epochs, model, criterion, optimizer, scheduler,
        train_loader, val_loader, device):
    trainer = create_supervised_trainer(model, optimizer, criterion, device=device)
    evaluator = create_supervised_evaluator(
        model,
        metrics={'accuracy': Accuracy(), 'nll': Loss(criterion)},
        device=device
    )

    RunningAverage(output_transform=lambda x: x).attach(trainer, 'loss')

    pbar = ProgressBar(persist=True)
    pbar.attach(trainer, metric_names='all')

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

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

        pbar.n = pbar.last_print_n = 0

    trainer.run(train_loader, max_epochs=epochs)

def main():
    epochs = 10
    train_loader, val_loader = get_loaders()
    model = get_model()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    optimizer, scheduler = get_optimizer(model)
    criterion = get_criterion()

    run(
        epochs=epochs,
        model=model,
        criterion=criterion,
        optimizer=optimizer,
        scheduler=scheduler,
        train_loader=train_loader,
        val_loader=val_loader,
        device=device
    )

2.2.4 Lightning

Lightning の場合、LightningModule を継承したクラス(Trainer クラス的なもの)を定義する必要があります。

各step(ex. training_step)の名前は決まっており、各 step を自分で埋めていきます。
また学習の実行自体は pytorch_lightning.Trainer クラスによって行われ、GPU や MixedPrecision、gradient accumulation 等の設定はこのクラスで設定します。
また metrics については Lightning 内には用意されていないため、自分で記載する必要があります。

import torch
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from share_funcs import get_model, get_loaders, get_criterion, get_optimizer

class MyLightninModule(pl.LightningModule):
    def __init__(self, num_class):
        super(MyLightninModule, self).__init__()
        self.model = get_model(num_class=num_class)
        self.criterion = get_criterion()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        loss = self.criterion(y_hat, y)
        logs = {'train_loss': loss}
        return {'loss': loss, 'log': logs, 'progress_bar': logs}

    def validation_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        preds = torch.argmax(y_hat, dim=1)
        return {'val_loss': self.criterion(y_hat, y), 'correct': (preds == y).float()}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        acc = torch.cat([x['correct'] for x in outputs]).mean()
        logs = {'val_loss': avg_loss, 'val_acc': acc}
        return {'avg_val_loss': avg_loss, 'log': logs}

    def configure_optimizers(self):
        # REQUIRED
        optimizer, scheduler = get_optimizer(model=self.model)
        return [optimizer], [scheduler]

    @pl.data_loader
    def train_dataloader(self):
        # REQUIRED
        return get_loaders()[0]

    @pl.data_loader
    def val_dataloader(self):
        # OPTIONAL
        return get_loaders()[1]


def main():
    epochs = 5
    num_class = 10
    output_path = './output/lightning'

    model = MyLightninModule(num_class=num_class)

    # most basic trainer, uses good defaults
    trainer = Trainer(
        max_nb_epochs=epochs,
        default_save_path=output_path,
        gpus=[0],
        # use_amp=False,
    )
    trainer.fit(model)

2.3 各フレームワークで実行したときのコンソール画面とアウトプット

2.3.1 デフォルト

コンソール画面

$ python train_default.py
Files already downloaded and verified
Files already downloaded and verified
Train start !
epoch 0 start !
100%|_____| 196/196 [00:05<00:00, 33.44it/s]
100%|_____| 40/40 [00:00<00:00, 50.43it/s]
epoch: 0  {'train_loss': 1.3714478426441854} {'valid_loss': 0.992230711877346, 'val_acc': tensor(0, device='cuda:0')}

アウトプット

なし

2.3.1 Catalyst

コンソール画面

$ python train_catalyst.py
1/5 * Epoch (train): 100% 196/196 [00:06<00:00, 30.09it/s, accuracy01=61.250, loss=1.058]
1/5 * Epoch (valid): 100% 40/40 [00:00<00:00, 49.75it/s, accuracy01=56.250, loss=1.053]
[2019-12-14 08:47:33,819]
1/5 * Epoch 1 (train): _base/lr=0.0010 | _base/momentum=0.9000 | _timers/_fps=58330.0450 | _timers/batch_time=0.0071 | _timers/data_time=0.0045 | _timers/model_time=0.0026 | accuracy01=52.0863 | loss=1.3634
1/5 * Epoch 1 (valid): _base/lr=0.0010 | _base/momentum=0.9000 | _timers/_fps=77983.3850 | _timers/batch_time=0.0146 | _timers/data_time=0.0126 | _timers/model_time=0.0019 | accuracy01=65.6250 | loss=0.9848
2/5 * Epoch (train): 100% 196/196 [00:06<00:00, 30.28it/s, accuracy01=63.750, loss=0.951]

アウトプット

• Tensorboard 等はデフォルトで出力されます
• weight も保存されます
• デフォルトで code まで残してくれるのはちょっとうれしいですね

catalyst
├── checkpoints
│   └── train.1.exception_KeyboardInterrupt.pth
├── code
│   ├── share_funcs.py
│   ├── train_catalyst.py
│   ├── train_default.py
│   └── train_lightning.py
├── log.txt
└── train_log
    └── events.out.tfevents.1576306176.FujimotoMac.local.41575.0

2.3.2 Ignite

コンソール画面

Catalyst よりもややスッキリした画面です。

$ python train_ignite.py
Epoch [1/10]: [196/196] 100%|________________, loss=1.14 [00:05<00:00]
Training Results - Epoch: 1  Avg accuracy: 0.69 Avg loss: 0.88
Validation Results - Epoch: 1  Avg accuracy: 0.65 Avg loss: 0.98
Epoch [2/10]: [196/196] 100%|________________, loss=0.813 [00:05<00:00]
Training Results - Epoch: 2  Avg accuracy: 0.78 Avg loss: 0.65
Validation Results - Epoch: 2  Avg accuracy: 0.70 Avg loss: 0.83

アウトプット

• なし
• 自分で保存する部分を書くか、Ignite 内のクラスを使う必要がありそうです

2.3.3 Lightning

コンソール画面

Lightning ではデフォルトでは tqdm 内のバーに全て表示するようです。

$ python train_lightning.py
Epoch 1: 100%|_____________| 236/236 [00:07<00:00, 30.75batch/s, batch_nb=195, gpu=0, loss=1.101, train_loss=1.06, v_nb=5]
Epoch 4:  41%|_____________| 96/236 [00:03<00:04, 32.28batch/s, batch_nb=95, gpu=0, loss=0.535, train_loss=0.524, v_nb=5]

アウトプット

• Lightning はディレクトリが重複した場合に version_x のように次のディレクトリを作って保存します。(それがかえって邪魔な場合もあり自分で checkpoint を定義することもありますが)
• Lightning の場合、meta_tags.csv に LightningModule に渡したパラメータが自動で保存されます
• Tensorboard 用の log もデフォルトで作成されます
• weight も checkpoints 内に保存されます
  • デフォルトでは各 epoch 毎に _ckpt_epoch_X.ckpt が作成され、古い epoch の ckpt を削除しているようです

lightning
└── lightning_logs
    ├── version_0
    │   └── checkpoints
    │       └── _ckpt_epoch_4.ckpt
    │   ├── media
    │   ├── meta.experiment
    │   ├── meta_tags.csv
    │   ├── metrics.csv
    │   └── tf
    │       └── events.out.tfevents.1576305970
    ├── version_1
    │   └── checkpoints
    │       └── _ckpt_epoch_3.ckpt
    │   ├── ...

2.4 その他めぼしいところ

いずれも Early Stopping 等は対応しています。

2.4.1 Catalyst

• catalyst.utils.set_global_seed() 等の再現性周りの関数も用意されている
• Dataset をより簡略化して書けるような関数もサポートされている
  • create_dataset, create_dataframe, prepare_dataset_labeling, split_dataframe
  • catalyst.utils.pandas
• Multi GPU や FP16 もサポートされている
• 公式 Tutorial のクオリティが高い
• 公式に Docker ファイルも置いてあり、インフラ周りの構成管理も意識したフレームワークを目指してるっぽい(多分)

2.4.2 Ignite

• Tensorboard や Logger も Ignite 内にあり、呼び出して使える
• 自由度が最も高そう
  • イベントの挟み方は慣れが必要そうだけど
• 公式リポジトリの下に置いてある

2.4.3 Lightning

• Multi GPU や FP16 もサポートされている

2.5 おすすめするとしたら

いずれもポテンシャルはあるため強制ではないです。下記は個人の感想です。

• 画像系コンペ初めてで、何からやればよいかよくわからない
  • → PyTorch Catalyst
• 画像系コンペは慣れきってて、殺意(金メダルを取りにいく強い気持ち)を持ってコンペに参加したい
  • → PyTorch Lightning か PyTorch Catalyst
• Classification・Segmentation に限らず色んな画像系タスクを取り組みたい
  • → PyTorch Lightning か PyTorch Catalyst
    • Catalyst 内には強化学習用のサンプルコードもある
• オサレに書きたい
  • → PyTorch Ignite
• PyTorch 公式のお膝元で安心してフレームワークを使いたい
  • → PyTorch Ignite
  • Ignite は 公式 PyTorch のリポジトリに置いてある

3. Catalyst・Ignite・Lightning を自由に行き来するために

ここで使うフレームワークを絞ってしまってもよいのですが、そもそも各フレームワークを行き来しやすいように書いていれば困らないはずです。 ですので PyTorch のコードを書き散らす上で意識しておくと良さそうなことをここにまとめます。

• ループの中身はなるべく取り出しておく
  • 各 step(ex. train 内の 1バッチごとの処理)の処理は抜き出せるように意識しておくと良さそうです
  • 少なくとも三重ループまで書き始めたら、ループの中身を抜き出せないか意識すると良さそうです
• 関数の引数を増やしすぎない
  • Class にしてインスタンス変数を使ったり、Config を一つにまとめて使っても良いと思います
• optimizer や model を呼び出す関数を作る
  • 個人の所感です
• Config を一つにまとめる
  • コンペに参加するときはなるべく一箇所に設定をまとめた方が管理が楽
  • 例
    • Config クラスを作る
    • Addict 等で呼び出しやすい辞書みたいなものを作る
    • YAML ファイルで書いた設定を読み出す

4. さいごに

本記事では PyTorch Catalyst・Ignite・Lightning の比較を行いました。
いずれも定型文を無くしたいという部分は一致していますが、それぞれ個性が出る結果となりました。
どのフレームワークもポテンシャルはあるため、もし触ってみて自分に合っていると思ったらコンペに出て使い倒してみると良いかと思います。

良い Kaggle(with PyTorch) ライフを！