/torchscope

Model Scope in PyTorch (include Params, FLOPs, Madds).

Primary LanguagePythonApache License 2.0Apache-2.0

torchscope

This is a neat plugin for scoping model in PyTorch. It is mainly based on the pytorch-summary and torchstat.

Installation

  • Install via pip
$ pip install torchscope
  • Install from source
$ pip install --upgrade git+https://github.com/Tramac/torchscope.git

Usage

from torchvision.models import resnet18
from torchscope import scope

model = resnet18()
scope(model, input_size=(3, 224, 224))
------------------------------------------------------------------------------------------------------
        Layer (type)               Output Shape          Params           FLOPs           Madds
======================================================================================================
            Conv2d-1          [1, 64, 112, 112]           9,408     118,013,952     235,225,088
       BatchNorm2d-2          [1, 64, 112, 112]             128       1,605,632       3,211,264
              ReLU-3          [1, 64, 112, 112]               0         802,816         802,816
         MaxPool2d-4            [1, 64, 56, 56]               0         802,816       1,605,632
            Conv2d-5            [1, 64, 56, 56]          36,864     115,605,504     231,010,304
       BatchNorm2d-6            [1, 64, 56, 56]             128         401,408         802,816
              ReLU-7            [1, 64, 56, 56]               0         200,704         200,704
            Conv2d-8            [1, 64, 56, 56]          36,864     115,605,504     231,010,304
       BatchNorm2d-9            [1, 64, 56, 56]             128         401,408         802,816
             ReLU-10            [1, 64, 56, 56]               0         200,704         200,704
           Conv2d-11            [1, 64, 56, 56]          36,864     115,605,504     231,010,304
      BatchNorm2d-12            [1, 64, 56, 56]             128         401,408         802,816
             ReLU-13            [1, 64, 56, 56]               0         200,704         200,704
           Conv2d-14            [1, 64, 56, 56]          36,864     115,605,504     231,010,304
      BatchNorm2d-15            [1, 64, 56, 56]             128         401,408         802,816
             ReLU-16            [1, 64, 56, 56]               0         200,704         200,704
           Conv2d-17           [1, 128, 28, 28]          73,728      57,802,752     115,505,152
      BatchNorm2d-18           [1, 128, 28, 28]             256         200,704         401,408
             ReLU-19           [1, 128, 28, 28]               0         100,352         100,352
           Conv2d-20           [1, 128, 28, 28]         147,456     115,605,504     231,110,656
      BatchNorm2d-21           [1, 128, 28, 28]             256         200,704         401,408
           Conv2d-22           [1, 128, 28, 28]           8,192       6,422,528      12,744,704
      BatchNorm2d-23           [1, 128, 28, 28]             256         200,704         401,408
             ReLU-24           [1, 128, 28, 28]               0         100,352         100,352
           Conv2d-25           [1, 128, 28, 28]         147,456     115,605,504     231,110,656
      BatchNorm2d-26           [1, 128, 28, 28]             256         200,704         401,408
             ReLU-27           [1, 128, 28, 28]               0         100,352         100,352
           Conv2d-28           [1, 128, 28, 28]         147,456     115,605,504     231,110,656
      BatchNorm2d-29           [1, 128, 28, 28]             256         200,704         401,408
             ReLU-30           [1, 128, 28, 28]               0         100,352         100,352
           Conv2d-31           [1, 256, 14, 14]         294,912      57,802,752     115,555,328
      BatchNorm2d-32           [1, 256, 14, 14]             512         100,352         200,704
             ReLU-33           [1, 256, 14, 14]               0          50,176          50,176
           Conv2d-34           [1, 256, 14, 14]         589,824     115,605,504     231,160,832
      BatchNorm2d-35           [1, 256, 14, 14]             512         100,352         200,704
           Conv2d-36           [1, 256, 14, 14]          32,768       6,422,528      12,794,880
      BatchNorm2d-37           [1, 256, 14, 14]             512         100,352         200,704
             ReLU-38           [1, 256, 14, 14]               0          50,176          50,176
           Conv2d-39           [1, 256, 14, 14]         589,824     115,605,504     231,160,832
      BatchNorm2d-40           [1, 256, 14, 14]             512         100,352         200,704
             ReLU-41           [1, 256, 14, 14]               0          50,176          50,176
           Conv2d-42           [1, 256, 14, 14]         589,824     115,605,504     231,160,832
      BatchNorm2d-43           [1, 256, 14, 14]             512         100,352         200,704
             ReLU-44           [1, 256, 14, 14]               0          50,176          50,176
           Conv2d-45             [1, 512, 7, 7]       1,179,648      57,802,752     115,580,416
      BatchNorm2d-46             [1, 512, 7, 7]           1,024          50,176         100,352
             ReLU-47             [1, 512, 7, 7]               0          25,088          25,088
           Conv2d-48             [1, 512, 7, 7]       2,359,296     115,605,504     231,185,920
      BatchNorm2d-49             [1, 512, 7, 7]           1,024          50,176         100,352
           Conv2d-50             [1, 512, 7, 7]         131,072       6,422,528      12,819,968
      BatchNorm2d-51             [1, 512, 7, 7]           1,024          50,176         100,352
             ReLU-52             [1, 512, 7, 7]               0          25,088          25,088
           Conv2d-53             [1, 512, 7, 7]       2,359,296     115,605,504     231,185,920
      BatchNorm2d-54             [1, 512, 7, 7]           1,024          50,176         100,352
             ReLU-55             [1, 512, 7, 7]               0          25,088          25,088
           Conv2d-56             [1, 512, 7, 7]       2,359,296     115,605,504     231,185,920
      BatchNorm2d-57             [1, 512, 7, 7]           1,024          50,176         100,352
             ReLU-58             [1, 512, 7, 7]               0          25,088          25,088
        AvgPool2d-59             [1, 512, 1, 1]               0          25,088          25,088
           Linear-60                  [1, 1000]         513,000         512,000       1,023,000
======================================================================================================
Total params: 11,689,512
Trainable params: 11,689,512
Non-trainable params: 0
Total FLOPs: 1,822,176,768
Total Madds: 3,639,535,640
----------------------------------------------------------------
Input size (MB): 0.14
Forward/backward pass size (MB): 14.26
Params size (MB): 11.15
Estimated Total Size (MB): 25.55
FLOPs size (GB): 1.82
Madds size (GB): 3.64
----------------------------------------------------------------

Note

This plugin only supports the following operations:

  • Conv2d
  • BatchNorm2d
  • Pool2d
  • ReLU
  • Upsample

Reference