torchsummaryよりtorchinfoがいいよという話

# なにこれ
- torchsummaryとtorch-summaryの話
- 結論：torchsummaryを使っていた人はtorchinfoに変えよう。

# 追記(2021.2.5)
- 名前がtorch-summaryからtorchinfoに変わりました。
- タイトル、結論、記事末尾のリンクだけ修正しました。

# 環境
tensorflow: 2.3.0
pytorch: 1.6.0
python: 3.7.2
torchsummary: 1.5.1
torch-summary: 1.4.1 

# summaryがほしいよね
書いたモデルをデバグする際に、さっと可視化できると非常に便利ですが、PyTorchにはtf.kerasの`model.summary()`がなく、`print(model)`することになります。

```python:keras_summary.py
# https://keras.io/api/models/model/

import tensorflow as tf

inputs = tf.keras.Input(shape=(3,))
x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs)
outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)


if __name__ == "__main__":
    model.summary()
```

```bash:output
$ python keras_summary.py

Model: "functional_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
dense (Dense)                (None, 4)                 16        
_________________________________________________________________
dense_1 (Dense)              (None, 5)                 25        
=================================================================
Total params: 41
Trainable params: 41
Non-trainable params: 0
_________________________________________________________________
```

```python:pytorch_print.py

import torch.nn as nn


model = nn.Sequential(
    nn.Linear(3, 4),
    nn.ReLU(),
    nn.Linear(4, 5),
    nn.Softmax(dim=1)
)


if __name__ == "__main__":
    print(model)
```

```bash:output

$ python pytorch_print.py

Model(
  (dense1): Linear(in_features=3, out_features=4, bias=True)
  (relu): ReLU()
  (dense2): Linear(in_features=4, out_features=5, bias=True)
  (softmax): Softmax(dim=1)
)
```

# そこで現れたtorchsummary

[issue](https://github.com/pytorch/pytorch/issues/2001)の流れで颯爽と現れ、この味気ないモデル表示をkerasっぽくして情報をリッチにできるのがtorchsummaryです。

```bash:installation
$ pip install torchsummary
```

```python:pytorch_summary.py

import torch.nn as nn
from torchsummary import summary


model = nn.Sequential(
    nn.Linear(3, 4),
    nn.ReLU(),
    nn.Linear(4, 5),
    nn.Softmax(dim=1)
)


if __name__ == "__main__":
    model = model.to("cuda")
    summary(model, (1,3))
```

```bash:output

$ python pytorch_summary.py

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Linear-1                 [-1, 1, 4]              16
              ReLU-2                 [-1, 1, 4]               0
            Linear-3                 [-1, 1, 5]              25
           Softmax-4                 [-1, 1, 5]               0
================================================================
Total params: 41
Trainable params: 41
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
----------------------------------------------------------------
```

しかしこのtorchsummary、開発が止まっている模様。
pypiからインストールするとコードが古く、[これ](https://github.com/sksq96/pytorch-summary/pull/46/commits/f9c420e46d248026136f033905a8a2507b196a6c)をしないとmultiple inputsに対応できませんでした。

# torch-summaryが更に情報をリッチに

- torchsummaryがmodelをユーザーがto("cuda")しなければならなかった点を解消
- 実際のコードを書き換える必要がない
- 親子関係が見やすくなる


```bash:installation
# torchsummaryが機能しなくなります。
$ pip install torch-summary
```

```bash:pytorch_torch-summry
# pipでインストールするだけで出力が変わる
$ python pytorch_summary.py

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
├─Sequential: 1-1                        [-1, 1, 5]                --
|    └─Linear: 2-1                       [-1, 1, 4]                16
|    └─ReLU: 2-2                         [-1, 1, 4]                --
|    └─Linear: 2-3                       [-1, 1, 5]                25
|    └─Softmax: 2-4                      [-1, 1, 5]                --
==========================================================================================
Total params: 41
Trainable params: 41
Non-trainable params: 0
Total mult-adds (M): 0.00
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================
```

この他にも、複数インプットがあるときに構造が見やすくなります。

```python:torchsummary_multiple.py

import torch.nn as nn
from torchsummary import summary


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(3, 4),
            nn.ReLU(),
            nn.Linear(4, 5),
            nn.Softmax(dim=1)
        )

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


if __name__ == "__main__":
    model = Model().cuda()
    summary(model, [(1, 3), (1, 3)])
```

```bash:output
$ python torchsummary_multiple.py

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Linear-1                 [-1, 1, 4]              16
              ReLU-2                 [-1, 1, 4]               0
            Linear-3                 [-1, 1, 5]              25
           Softmax-4                 [-1, 1, 5]               0
            Linear-5                 [-1, 1, 4]              16
              ReLU-6                 [-1, 1, 4]               0
            Linear-7                 [-1, 1, 5]              25
           Softmax-8                 [-1, 1, 5]               0
================================================================
```

```python:torch-summary_multiple.py

import torch.nn as nn
from torchsummary import summary


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(3, 4),
            nn.ReLU(),
            nn.Linear(4, 5),
            nn.Softmax(dim=1)
        )

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

if __name__ == "__main__":
    model = Model()
    summary(model, [(1, 3), (1, 3)])
```

```bash:output
$ python torch-summary_multiple.py

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
├─Sequential: 1-1                        [-1, 1, 5]                --
|    └─Linear: 2-1                       [-1, 1, 4]                16
|    └─ReLU: 2-2                         [-1, 1, 4]                --
|    └─Linear: 2-3                       [-1, 1, 5]                25
|    └─Softmax: 2-4                      [-1, 1, 5]                --
├─Sequential: 1-2                        [-1, 1, 5]                (recursive)
|    └─Linear: 2-5                       [-1, 1, 4]                (recursive)
|    └─ReLU: 2-6                         [-1, 1, 4]                --
|    └─Linear: 2-7                       [-1, 1, 5]                (recursive)
|    └─Softmax: 2-8                      [-1, 1, 5]                --
==========================================================================================
```

# 実はtorchsummaryXもある

torch-summaryは、torchsummaryとtorchsummaryXの後継と自称しています。
torchsummaryXから乗り換える際は下の2つの変更が必要かと思います。
- `from torchsummaryX` => `from torchsummary`
- summary()の引数`input`は、`input.shape`に変更

# まとめ
- (オススメ) [torchinfo(GitHub)](https://github.com/TylerYep/torchinfo)
- [torchsummary(GitHub)](https://github.com/sksq96/pytorch-summary)
- [torchsummaryX(GitHub)](https://github.com/nmhkahn/torchsummaryX)