ChainerのMNISTのコードをPyTorch+Igniteで書き直してみた

# TL;DR
ChainerのMNISTのコードをPyTorchに書き直して見た感想は、ほとんど違いがなかったです。
違いがあった部分は、Chainerで言うUpdaterより上位層、PyTorchで言うIgnite層でした。
しかも、実は[chainer-pytorch-migration][3]を使うと、Chainerで使っていたExtentionsをIgniteでも使えたりして、かなりChainerライクにPyTorch+Igniteを使うことができます。
今までChainerを使っていた方も自然とPyTorch+Igniteに慣れることができるのではと思います。

# PFNによるChainerの開発停止とPyTorch採用

ここ[\[1\]][1]に記載されているように、PFNはChainerの開発を終了し、PyTorchに移行することが発表されました。

> 株式会社Preferred Networks（本社：東京都千代田区、代表取締役社長：西川徹、プリファードネットワークス、以下、PFN）は、研究開発の基盤技術である深層学習フレームワークを、自社開発のChainer™から、PyTorchに順次移行します。同時に、PyTorchを開発する米FacebookおよびPyTorchの開発者コミュニティと連携し、PyTorchの開発に参加します。なお、Chainerは、本日公開されたメジャーバージョンアップとなる最新版v7をもってメンテナンスフェーズに移行します。Chainerユーザー向けには、PyTorchへの移行を支援するドキュメントおよびライブラリを提供します。

今すぐにChainerを使えなくなるわけではないですが、Chainerユーザは徐々に他のフレームワークに移行せざるを得ない状況になっています。

# PFNによるChainerからPyTorch移行のサポート
Chainer開発終了は急な発表で多くのユーザは困惑したと思いますが、PFNもその状況を想定してPyTorchへの移行を支援する[ドキュメント \[2\]][2]および[ライブラリ \[3\]][3]を提供しています。

上記、ドキュメントを見ると、ChainerとPyTorch+Igniteの対応関係は以下のようです。

<img width="758" alt="スクリーンショット 2019-12-17 15.32.28.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/111482/b86cbf5a-5450-965d-a32b-a2d8331ab25c.png"> cited from [\[2\]][2]

上記からわかることは、

- ChainerのOptimizerまでの役割は、PyTorchが対応している
- ChainerのUpdater・Trainerの役割は、Igniteが対応している

なので、学習ステップを自分で書く場合はPyTorchのみで書けるが、ChainerのTrainerのように学習ステップもフレームワークで対応してもらいたい場合はPyTorch+Igniteを使う必要があることがわかります。

# さっそくChainerからPyTorch+Igniteへ移行してみた
## 移行対象となるコード
下記リンク先にChainerのTrainerを使って、MNISTの訓練・推論を行うNotebookがあります。

- [Chainer Begginer's Hands-on » Trainerを使ってみよう](https://chainer-colab-notebook.readthedocs.io/ja/latest/notebook/hands_on/chainer/begginers_hands_on/12_Try_Trainer_class.html)

今回は上記コードをPyTorch+Igniteを使って書き直してみたいと思います。

## 各ステップの移行方法
### サンプルデータセットの読み込み部分

```python
from chainer.datasets import mnist

train, test = mnist.get_mnist()
```
↓

```python
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

data_transform = ToTensor()

train = MNIST(download=True, root=".", transform=data_transform, train=True)
test = MNIST(download=False, root=".", transform=data_transform, train=False)
```

- 違い
    - ここはサンプルデータセットの読み込み部分なのでそれなりに違いがある部分だと思いますが、実用途では違いがあってもあまり関係がない部分です。

### Iterator -> DataLoader

```python
from chainer import iterators

batchsize = 128

train_iter = iterators.SerialIterator(train, batchsize)
test_iter = iterators.SerialIterator(test, batchsize, False, False)
```
↓

```python
from torch.utils.data import DataLoader

batch_size = 128

train_loader = DataLoader(train, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test, batch_size=batch_size, shuffle=False)
```

- 違い（ほとんど同じですね）
    - 引数がちょっと違うくらいですかね。

### モデルの準備

```python

import chainer
import chainer.links as L
import chainer.functions as F

class MLP(chainer.Chain):

    def __init__(self, n_mid_units=100, n_out=10):
        super(MLP, self).__init__()
        with self.init_scope():
            self.l1=L.Linear(None, n_mid_units)
            self.l2=L.Linear(None, n_mid_units)
            self.l3=L.Linear(None, n_out)


    def forward(self, x):
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        return self.l3(h2)

gpu_id = 0  # Set to -1 if you don't have a GPU

model = L.Classifier(model)
if gpu_id >= 0:
    model.to_gpu(gpu_id)
```
↓

```python
from torch import nn
import torch.nn.functional as F
import torch

class MLP(nn.Module):

    def __init__(self, n_mid_units=100, n_out=10):
        super(MLP, self).__init__()
        self.l1 = nn.Linear(784, n_mid_units)
        self.l2 = nn.Linear(n_mid_units, n_mid_units)
        self.l3 = nn.Linear(n_mid_units, n_out)

    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        h3 = self.l3(h2)
        return F.log_softmax(h3, dim=1)

device = 'cuda:0'

model = MLP()
```

- 違い（ほとんど同じですね）
    - PyTorchの場合は、`Linear`の`in_features`を`None`として省略することができないようです。
    - PyTorchの場合は、Chainerの`L.Classifier`が存在しないので、最終層で`F.log_softmax(h3, dim=1)`を明示的に計算している。
    - (フレームワークの違いというよりはデータセットの形式の違い)PyTorchのMNISTのデータは2次元なので、`x = torch.flatten(x, start_dim=1)`として1次元にしている。

### Optimizerの準備

```python
from chainer import optimizers

lr = 0.01

optimizer = optimizers.SGD(lr=lr)
optimizer.setup(model)
```
↓

```python
from torch import optim

lr = 0.01

# 最適化手法の選択
optimizer = optim.SGD(model.parameters(), lr=lr)
```

- 違い（ほとんど同じですね）
    - 引数がちょっと違うくらいですかね。

### Updater -> Ignite

```python
from chainer import training

updater = training.StandardUpdater(train_iter, optimizer, device=gpu_id)
```
↓

```python
from ignite.engine import create_supervised_trainer

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

- 違い（ほとんど同じですね）
    - 引数がちょっと違うくらいですかね。

### 拡張機能の追加

```python
from chainer.training import extensions

trainer = training.Trainer(
    updater, (max_epoch, 'epoch'), out='mnist_result'
)

trainer.extend(extensions.LogReport())
.
.
.
trainer.extend(extensions.Evaluator(test_iter, model, device=gpu_id))
```
↓

```python
from ignite.engine import create_supervised_evaluator
from ignite.metrics import Accuracy, Loss
from ignite.engine import Events

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

training_history = {'accuracy':[],'loss':[]}
validation_history = {'accuracy':[],'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']
    training_history['accuracy'].append(avg_accuracy)
    training_history['loss'].append(avg_nll)
    print(
        "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(test_loader)
    metrics = evaluator.state.metrics
    avg_accuracy = metrics['accuracy']
    avg_nll = metrics['nll']
    validation_history['accuracy'].append(avg_accuracy)
    validation_history['loss'].append(avg_nll)
    print(
        "Validation Results - Epoch: {}  Avg accuracy: {:.2f} Avg loss: {:.2f}"
        .format(engine.state.epoch, avg_accuracy, avg_nll))

# Create snapshot
from ignite.handlers import ModelCheckpoint

checkpointer = ModelCheckpoint(
    './models',
    'MNIST',
    save_interval=1,
    n_saved=2, 
    create_dir=True, 
    save_as_state_dict=True,
    require_empty=False,
)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpointer, {'MNIST': model})
```

- 違い
    - ためしに、Igniteで以下をしてみたのですが、Chainerと書き方はだいぶ異なります。ただ、実装さえすれば今までできたことができないということはなさそうです。
        - 訓練データでの精度・ロスを表示・保存
        - 検証データでの精度・ロスを表示・保存
        - １epoch ごとにスナップショットを保存

### 訓練の実行

```python
trainer.run()
```

↓

```python
max_epochs = 10
trainer.run(train_loader, max_epochs=max_epochs)
```

- 違い（ほとんど同じですね）
    - 引数がちょっと違うくらいですかね。

### 移行後のコード
実際に移行してColaboratoryで動くNotebookは以下にあります。上記で説明したコードだけでなく、以下も含まれているのでよろしければお手元で実行してみてください。

- 訓練・検証データでの精度・ロスのプロットする
- スナップショットからモデルをロードし、推論する

https://drive.google.com/open?id=1NqHYJjFz-dl1tWP8kMO0y0kCZ9-ZWLxi

# おまけ
実は、[chainer-pytorch-migration][3]を使うと、Chainerで使っていたextensionsをIgniteでも使えます！ Chainerのextensionsが恋しい方はchainer-pytorch-migrationを使ってみてください。

```python
import chainer_pytorch_migration as cpm
import chainer_pytorch_migration.ignite
from chainer.training import extensions

optimizer.target = model
trainer.out = 'result'

cpm.ignite.add_trainer_extension(trainer, optimizer, extensions.LogReport())
cpm.ignite.add_trainer_extension(trainer, optimizer, extensions.ExponentialShift('lr', 0.9, 1.0, 0.1))
cpm.ignite.add_trainer_extension(trainer, optimizer, extensions.PrintReport(
    ['epoch', 'iteration', 'loss', 'lr']))

max_epochs = 10
trainer.run(train_loader, max_epochs=max_epochs)
```

# Reference
- [1] [Preferred Networks、深層学習の研究開発基盤をPyTorchに移行][1]
- [2] [Framework Migration Guide][2]
- [3] [github.com/chainer/chainer-pytorch-migration][3]

[1]:https://preferred.jp/ja/news/pr20191205/
[2]:https://chainer.github.io/migration-guide/
[3]:https://github.com/chainer/chainer-pytorch-migration