PyTorch - model summery

在 Keras 中，model.summary 是非常好用的功能，可以清楚的知道 CNN 每一層的 Kernal size 還有 output shape，然而在 PyTorch 中沒有類似功能，在搭建 CNN 層的過程，都要自己計算，蠻不方便的。
在 Keras 中，簡單到不需要特別說明，只有在練習過程中簡略帶過。沒想到這功能是如此的方便!!

在 Keras 中的模型搭建，利用 model.summary()即可 print 出漂亮的表格：

from keras import layers
from keras import models
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Dropout(0.5))  #加入dropout層，預防過擬合
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))
model.summary()

PyTorch 是否也能做到!? google 了一下，有人把 PyTorch 打印 model summary 的function 從 Keras 移植過來了!! 直接 Clone sksq96 大神的 GitHub repo. ： pytorch-summary 即可。

我的做法比較粗暴，我直接將 torchsummary.py 丟到我的 Conda 環境資料夾，就可以直接 import 使用。

使用方法很簡單。
將你的 model 和 input_size (即input shape) 輸入至 summary() 中即可。

from torchsummary import summary
summary(your_model, input_size=(channels, H, W))

以下用簡單的CNN範例:

import torch
import torch.nn as nn

# Create CNN Model
class CNN_Model(nn.Module):
    def __init__(self):
        super(CNN_Model, self).__init__()
        # Convolution 1 , input_shape=(1,28,28)
        self.cnn1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=0) #output_shape=(16,24,24)
        self.relu1 = nn.ReLU() # activation
        # Max pool 1
        self.maxpool1 = nn.MaxPool2d(kernel_size=2) #output_shape=(16,12,12)
        # Convolution 2
        self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=0) #output_shape=(32,8,8)
        self.relu2 = nn.ReLU() # activation
        # Max pool 2
        self.maxpool2 = nn.MaxPool2d(kernel_size=2) #output_shape=(32,4,4)
        # Fully connected 1 ,#input_shape=(32*4*4)
        self.fc1 = nn.Linear(32 * 4 * 4, 10) 
    
    def forward(self, x):
        # Convolution 1
        out = self.cnn1(x)
        out = self.relu1(out)
        # Max pool 1
        out = self.maxpool1(out)
        # Convolution 2 
        out = self.cnn2(out)
        out = self.relu2(out)
        # Max pool 2 
        out = self.maxpool2(out)
        out = out.view(out.size(0), -1)
        # Linear function (readout)
        out = self.fc1(out)
        return out

model = CNN_Model()

from torchsummary import summary
summary(model.cuda(), (1, 28, 28))

一目了然，這樣在建構 CNN model 時清楚多了!!

也私心希望 PyTorch 官方把這功能加進去…