[Pytorch] torch summary

모델 내 레이어에 따른 output shape, parameter 개수에 대한 정보를 표로 쉽게 볼 수 있도록 하는 파이썬 패키지

 

설치

아래와 같이 pip을 통해 설치 할 수 있다

pip install torchsummary

 

 

사용법

from torchsummary import summary

summary(model, input_size = (channels, H, W))

• input_size의 차원 수는 2/3/4차원 모두 가능
• model은 GPU(cuda)에 있어야 사용 가능

 

결과

• 각 Layer 별 Output Shape, Parameter 개수에 대한 정보가 출력되는 것을 볼 수 있다

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
         ZeroPad2d-1            [-1, 3, 81, 81]               0
Conv2dStaticSamePadding-2           [-1, 32, 40, 40]         864
       BatchNorm2d-3           [-1, 32, 40, 40]              64
MemoryEfficientSwish-4           [-1, 32, 40, 40]             0
         ZeroPad2d-5           [-1, 32, 42, 42]               0
Conv2dStaticSamePadding-6           [-1, 32, 40, 40]         288
       BatchNorm2d-7           [-1, 32, 40, 40]              64
MemoryEfficientSwish-8           [-1, 32, 40, 40]             0
          Identity-9             [-1, 32, 1, 1]               0
Conv2dStaticSamePadding-10              [-1, 8, 1, 1]        264
MemoryEfficientSwish-11              [-1, 8, 1, 1]            0
         Identity-12              [-1, 8, 1, 1]               0
Conv2dStaticSamePadding-13             [-1, 32, 1, 1]        288
         Identity-14           [-1, 32, 40, 40]               0
Conv2dStaticSamePadding-15           [-1, 16, 40, 40]        512
      BatchNorm2d-16           [-1, 16, 40, 40]              32
      MBConvBlock-17           [-1, 16, 40, 40]               0
         Identity-18           [-1, 16, 40, 40]               0
Conv2dStaticSamePadding-19           [-1, 96, 40, 40]        1,536
      BatchNorm2d-20           [-1, 96, 40, 40]              192
MemoryEfficientSwish-21           [-1, 96, 40, 40]            0
        ZeroPad2d-22           [-1, 96, 41, 41]               0
show more (open the raw output data in a text editor) ...

Input size (MB): 0.07
Forward/backward pass size (MB): 28.31
Params size (MB): 15.30
Estimated Total Size (MB): 43.69
----------------------------------------------------------------