This is a neat plugin for scoping model in PyTorch. It is mainly based on the pytorch-summary and torchstat.
- Install via pip
$ pip install torchscope
- Install from source
$ pip install --upgrade git+https://github.com/Tramac/torchscope.git
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
----------------------------------------------------------------
This plugin only supports the following operations:
- Conv2d
- BatchNorm2d
- Pool2d
- ReLU
- Upsample