[Pytorch] model summary : torchsummary, torchinfo로 layer별 output shape 확인하기

torchsummary 

# Convolution layer로 구성된 model

channel=1, 28x28 사이즈의 흑백 이미지를 input으로, convolution layer를 이용한 이미지의 noise를 제거하는 encoder, decoder 구조.

from torchsummary import summary
summary(model, input_size=(1,28, 28), batch_size=32) # channel수, w, h

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [32, 16, 28, 28]             160
              ReLU-2           [32, 16, 28, 28]               0
            Conv2d-3           [32, 16, 28, 28]           2,320
              ReLU-4           [32, 16, 28, 28]               0
        BasicBlock-5           [32, 16, 28, 28]               0
         AvgPool2d-6           [32, 16, 14, 14]               0
            Conv2d-7            [32, 8, 14, 14]           1,160
              ReLU-8            [32, 8, 14, 14]               0
            Conv2d-9            [32, 8, 14, 14]             584
             ReLU-10            [32, 8, 14, 14]               0
       BasicBlock-11            [32, 8, 14, 14]               0
        AvgPool2d-12              [32, 8, 7, 7]               0
          Encoder-13              [32, 8, 7, 7]               0
           Conv2d-14              [32, 8, 7, 7]             584
             ReLU-15              [32, 8, 7, 7]               0
           Conv2d-16              [32, 8, 7, 7]             584
             ReLU-17              [32, 8, 7, 7]               0
       BasicBlock-18              [32, 8, 7, 7]               0
  ConvTranspose2d-19            [32, 8, 14, 14]             264
           Conv2d-20           [32, 16, 14, 14]           1,168
             ReLU-21           [32, 16, 14, 14]               0
           Conv2d-22           [32, 16, 14, 14]           2,320
             ReLU-23           [32, 16, 14, 14]               0
       BasicBlock-24           [32, 16, 14, 14]               0
  ConvTranspose2d-25           [32, 16, 28, 28]           1,040
           Conv2d-26            [32, 1, 28, 28]             145
          Decoder-27            [32, 1, 28, 28]               0
================================================================
Total params: 10,329
Trainable params: 10,329
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.10
Forward/backward pass size (MB): 26.32
Params size (MB): 0.04
Estimated Total Size (MB): 26.45
----------------------------------------------------------------

 

torchinfo

# LSTM layer로 구성된 model

LSTM은 input, output shape가 복잡해서 그런지 tochsummary module을 사용했을 때 에러가 난다. 대신 사용했던 torchinfo 모듈의 summary

 

입력으로 2개의 token을 받고, BOW로 문자 --> 숫자로 encoding한 뒤 embedding layer, LSTM layer를 거쳐 next token 하나를 예측하는 구조.

 

batch size=64

BOW length=12148

LSTM에서 한번에 input하는 token 수 =2 (2개 toekn을 입력받고 next 1개 token을 예측함)

output 12148개 중 argmax 값이 next token

from torchinfo import summary
summary(model, input_size=(64, 2), dtypes=[torch.long])

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
LSTM                                     [64, 12148]               --
├─Embedding: 1-1                         [64, 2, 16]               194,368
├─LSTM: 1-2                              [64, 2, 64]               154,112
├─Linear: 1-3                            [64, 12148]               1,567,092
├─ReLU: 1-4                              [64, 12148]               --
├─Linear: 1-5                            [64, 12148]               147,586,052
==========================================================================================
Total params: 149,501,624
Trainable params: 149,501,624
Non-trainable params: 0
Total mult-adds (G): 9.58
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 12.52
Params size (MB): 598.01
Estimated Total Size (MB): 610.53
==========================================================================================

 

* dtype=[torch.long]을 입력해주지 않으면 embedding layer에서 아래와 같은 에러가 난다.

-> 2210 return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse) 
RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)

...

RuntimeError: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: []

이 오류 메시지는 모델의 Embedding 레이어가 입력으로 Long 또는 Int 타입의 텐서를 기대하지만, FloatTensor 타입의 텐서가 전달되었다는 뜻. torchsummary와는 달리 입력 텐서의 type을 지정할 수 있어서 dtype 파라미터에 적절한 입력 텐서의 type을 써주니 해결됐다.

 

torchinfo 공식 documentation

https://github.com/TylerYep/torchinfo