[torchsummary] Pytorch에서 keras처럼 모델 출력하기

[수정 2022/06/12] torchsummary 패키지는 torchinfo 라는 패키지로 업그레이드되며 이름이 바뀌었습니다.

 

pytorch와 함께 사용할 수 있는 torchsummary라는 파이썬 패키지입니다.

pip install torchsummary
git clone https://github.com/sksq96/pytorch-summary

위와 같이 pip를 통해서 설치할 수 있고, 아래와 같이 직접 깃허브에 있는 소스코드를 다운받아 사용할 수 있습니다.

 

Pytorch는 keras에 비해 모델을 출력했을 때 깔끔하게 볼 수 가 없습니다. keras의 경우는 파라미터의 개수 레이어의 종류, 넘버 등을 표로 정리해서 알려주는 것과 비교해보면 말이죠.

위 torchsummary 라는 패키지가 pytorch에서 keras와 같은 형태로 모델을 출력해줍니다.

 

사용법은 간단한데요,

from torchsummary import summary

your_model = CNN().cuda()
summary(your_model, input_size=(channels, H, W))

우선 torchsummary라는 패키지(파일)에서 summary를 임포트하고, 모델을 선언한 후에 인풋의 크기와 함께 summary 함수로 넣어주면 됩니다.

input_size 가 없으면 에러가 뜹니다. 그리고 사용해보며 알게 된 점으로는

• input size는 2차원이 아니라 3차원, 4차원 상관없다.
• 모델이 만약 input이 2개를 받는다면 저 summary 함수 안에 input을 2개 넣어주어도 된다.
• 만약 위와 같이 사용할 때에는 모델이 cuda(), 다시말해 GPU에 올라가있는게 아니면 에러가 난다.

GPU에 없을 때 왜 에러가 나는지 알기 위해 summary 코드를 읽어보았는데, 이 코드에서는 cuda(), cpu() 모두 있습니다. 즉 GPU, CPU 상관이 없습니다.

만약 GPU가 없을 때, 사용할 수 없을 때 모델을 cpu 상에 올려놓고 summary 를 출력하고 싶으시다면 아래와 같이 하셔야 합니다.

from torchsummary import summary

your_model = CNN().cpu()
summary(your_model, input_size=(channels, H, W), device='cpu')

쉽게 말해 device라는 인자에서 cpu라고 명시를 하지 않으시면 cuda로 작동하게 되어 오류가 발생합니다.

 

 

아래는 예시 코드와 그에 따른 결과입니다. torchsummary github에서 가져왔습니다.

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") # PyTorch v0.4.0
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
----------------------------------------------------------------

문제없이 잘 작동하는 것을 보실 수 있습니다.

저 input size, forward/backward pass size, 부터 total size를 어떻게 계산하는지 궁금하시면 직접 github에서 코드를 뜯어보시면 됩니다. 

한 번 본적이 있는데 단순히 모든 파라미터가 cuda 에서 float 변수일 때, 즉 4 byte/1개당 일때로 놓고 계산합니다.

 

저 estimated total size 의 결과와 torch.save()로 모델을 저장했을 때의 크기는 비교해보지 않아서 모르겠습니다. 궁금하네요. 해보신 분은 댓글로 알려주셔도 좋을 것 같습니다.

---

torchsummary는 제가 주로 사용하는 패키지 중 하나입니다. 만약 모델이 복잡하게 블럭들이 엉켜 있는 경우에는 잘 작동하지 않지만 그냥 전체적인 흐름을 볼 때 좋은 것 같습니다. 머신 러닝이나 pytorch를 처음 접하시는 분이라면 도움이 많이 될 것 같습니다.