torch-summary로 케라스 처럼 모델 요약(summary)하기

케라스를 사용하다보면 몇 가지 편리한 기능들이 있다. 파이토치가 좀더 파이썬 객체 지향적으로 개발을 할 수 있다는 장점도 있지만, 케라스 API로 제공되는 편리한 기능들이 빠져 있고 직접 구현을 해야하는 부분들이 있어서 불편한 점들도 있다. 텐서플로우v2 부터는 케라스가 통합되어 사용할 수 있는데, 파이토치를 쓰다보면 케라스의 고수준 API 들이 그리울 때가 있는 것 같다.

그중에서 한 가지를 뽑자면 아래처럼 모델을 summary 해주는 기능이다.

그런데 파이토치에서 케라스처럼 저렇게 모델을 summary 해주는 기능은 없을까?

pytorch에서 모델을 출력하게 되면 아래처럼 나오는데,

아무래도 이렇게 보는 것은 한눈에 와닿지가 않고 얼마나 많은 파라미터가 필요한지도 알 수가 없다.

pytorch-summary

파이토치 프레임워크에는 이런 기능이 기본적으로 제공되지 않기 때문에 불편하다고 생각했는데,

찾아보니 누가 유틸리티 라이브러리로 만들어 둔 개발자가 있었다.

sksq96/pytorch-summary: Model summary in PyTorch similar to model.summary() in Keras (github.com)

pytorch-summary 설치

pip install torchsummary

이것을 사용하면 keras의 summary처럼 거의 유사하게 확인해볼 수 있다.

예제

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
model = Net().to(device)

summary(model, (1, 28, 28))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------