Added ability to print model size and summary at run-time

[sarthakpati]

Fixes #318

Proposed Changes

• using torchinfo to generate model summary
• this should be extended so that it is parsed and then memory comparison is done to give meaningful error(s) at run-time

Checklist

• I have read the CONTRIBUTING guide
• My PR is based from the current GaNDLF master
• Non-breaking change (does not break existing functionality): provide as many details as possible for any breaking change
• Function/class source code documentation added/updated
• Code has been blacked for style consistency
• If applicable, version information has been updated in GANDLF/version.py
• If adding a git submodule, add to list of exceptions for black styling in pyproject.toml file
• Usage documentation has been updated, if appropriate
• History has been updated, if appropriate
• Tests added or modified to cover the changes; if coverage is reduced, please give explanation
• If customized dependency installation is required (i.e., a separate pip install step is needed for PR to be functional), please ensure it is reflected in all the files that control the CI, namely: python-test.yml, and all docker files [1,2,3,4]

[sarthakpati]

The output is something similar to the following:

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
DenseNet                                 [1, 3]                    --
├─Sequential: 1-1                        [1, 1024, 8, 4]           --
│    └─Conv2d: 2-1                       [1, 64, 128, 64]          9,408
│    └─BatchNorm2d: 2-2                  [1, 64, 128, 64]          128
│    └─ReLU: 2-3                         [1, 64, 128, 64]          --
│    └─MaxPool2d: 2-4                    [1, 64, 64, 32]           --
│    └─_DenseBlock: 2-5                  [1, 256, 64, 32]          --
│    │    └─_DenseLayer: 3-1             [1, 96, 64, 32]           45,440
│    │    └─_DenseLayer: 3-2             [1, 128, 64, 32]          49,600
│    │    └─_DenseLayer: 3-3             [1, 160, 64, 32]          53,760
│    │    └─_DenseLayer: 3-4             [1, 192, 64, 32]          57,920
│    │    └─_DenseLayer: 3-5             [1, 224, 64, 32]          62,080
│    │    └─_DenseLayer: 3-6             [1, 256, 64, 32]          66,240
│    └─_Transition: 2-6                  [1, 128, 32, 16]          --
│    │    └─BatchNorm2d: 3-7             [1, 256, 64, 32]          512
│    │    └─ReLU: 3-8                    [1, 256, 64, 32]          --
│    │    └─Conv2d: 3-9                  [1, 128, 64, 32]          32,768
│    │    └─AvgPool2d: 3-10              [1, 128, 32, 16]          --
│    └─_DenseBlock: 2-7                  [1, 512, 32, 16]          --
│    │    └─_DenseLayer: 3-11            [1, 160, 32, 16]          53,760
│    │    └─_DenseLayer: 3-12            [1, 192, 32, 16]          57,920
│    │    └─_DenseLayer: 3-13            [1, 224, 32, 16]          62,080
│    │    └─_DenseLayer: 3-14            [1, 256, 32, 16]          66,240
│    │    └─_DenseLayer: 3-15            [1, 288, 32, 16]          70,400
│    │    └─_DenseLayer: 3-16            [1, 320, 32, 16]          74,560
│    │    └─_DenseLayer: 3-17            [1, 352, 32, 16]          78,720
│    │    └─_DenseLayer: 3-18            [1, 384, 32, 16]          82,880
│    │    └─_DenseLayer: 3-19            [1, 416, 32, 16]          87,040
│    │    └─_DenseLayer: 3-20            [1, 448, 32, 16]          91,200
│    │    └─_DenseLayer: 3-21            [1, 480, 32, 16]          95,360
│    │    └─_DenseLayer: 3-22            [1, 512, 32, 16]          99,520
│    └─_Transition: 2-8                  [1, 256, 16, 8]           --
│    │    └─BatchNorm2d: 3-23            [1, 512, 32, 16]          1,024
│    │    └─ReLU: 3-24                   [1, 512, 32, 16]          --
│    │    └─Conv2d: 3-25                 [1, 256, 32, 16]          131,072
│    │    └─AvgPool2d: 3-26              [1, 256, 16, 8]           --
│    └─_DenseBlock: 2-9                  [1, 1024, 16, 8]          --
│    │    └─_DenseLayer: 3-27            [1, 288, 16, 8]           70,400
│    │    └─_DenseLayer: 3-28            [1, 320, 16, 8]           74,560
│    │    └─_DenseLayer: 3-29            [1, 352, 16, 8]           78,720
│    │    └─_DenseLayer: 3-30            [1, 384, 16, 8]           82,880
│    │    └─_DenseLayer: 3-31            [1, 416, 16, 8]           87,040
│    │    └─_DenseLayer: 3-32            [1, 448, 16, 8]           91,200
│    │    └─_DenseLayer: 3-33            [1, 480, 16, 8]           95,360
│    │    └─_DenseLayer: 3-34            [1, 512, 16, 8]           99,520
│    │    └─_DenseLayer: 3-35            [1, 544, 16, 8]           103,680
│    │    └─_DenseLayer: 3-36            [1, 576, 16, 8]           107,840
│    │    └─_DenseLayer: 3-37            [1, 608, 16, 8]           112,000
│    │    └─_DenseLayer: 3-38            [1, 640, 16, 8]           116,160
│    │    └─_DenseLayer: 3-39            [1, 672, 16, 8]           120,320
│    │    └─_DenseLayer: 3-40            [1, 704, 16, 8]           124,480
│    │    └─_DenseLayer: 3-41            [1, 736, 16, 8]           128,640
│    │    └─_DenseLayer: 3-42            [1, 768, 16, 8]           132,800
│    │    └─_DenseLayer: 3-43            [1, 800, 16, 8]           136,960
│    │    └─_DenseLayer: 3-44            [1, 832, 16, 8]           141,120
│    │    └─_DenseLayer: 3-45            [1, 864, 16, 8]           145,280
│    │    └─_DenseLayer: 3-46            [1, 896, 16, 8]           149,440
│    │    └─_DenseLayer: 3-47            [1, 928, 16, 8]           153,600
│    │    └─_DenseLayer: 3-48            [1, 960, 16, 8]           157,760
│    │    └─_DenseLayer: 3-49            [1, 992, 16, 8]           161,920
│    │    └─_DenseLayer: 3-50            [1, 1024, 16, 8]          166,080
│    └─_Transition: 2-10                 [1, 512, 8, 4]            --
│    │    └─BatchNorm2d: 3-51            [1, 1024, 16, 8]          2,048
│    │    └─ReLU: 3-52                   [1, 1024, 16, 8]          --
│    │    └─Conv2d: 3-53                 [1, 512, 16, 8]           524,288
│    │    └─AvgPool2d: 3-54              [1, 512, 8, 4]            --
│    └─_DenseBlock: 2-11                 [1, 1024, 8, 4]           --
│    │    └─_DenseLayer: 3-55            [1, 544, 8, 4]            103,680
│    │    └─_DenseLayer: 3-56            [1, 576, 8, 4]            107,840
│    │    └─_DenseLayer: 3-57            [1, 608, 8, 4]            112,000
│    │    └─_DenseLayer: 3-58            [1, 640, 8, 4]            116,160
│    │    └─_DenseLayer: 3-59            [1, 672, 8, 4]            120,320
│    │    └─_DenseLayer: 3-60            [1, 704, 8, 4]            124,480
│    │    └─_DenseLayer: 3-61            [1, 736, 8, 4]            128,640
│    │    └─_DenseLayer: 3-62            [1, 768, 8, 4]            132,800
│    │    └─_DenseLayer: 3-63            [1, 800, 8, 4]            136,960
│    │    └─_DenseLayer: 3-64            [1, 832, 8, 4]            141,120
│    │    └─_DenseLayer: 3-65            [1, 864, 8, 4]            145,280
│    │    └─_DenseLayer: 3-66            [1, 896, 8, 4]            149,440
│    │    └─_DenseLayer: 3-67            [1, 928, 8, 4]            153,600
│    │    └─_DenseLayer: 3-68            [1, 960, 8, 4]            157,760
│    │    └─_DenseLayer: 3-69            [1, 992, 8, 4]            161,920
│    │    └─_DenseLayer: 3-70            [1, 1024, 8, 4]           166,080
│    └─BatchNorm2d: 2-12                 [1, 1024, 8, 4]           2,048
├─Linear: 1-2                            [1, 3]                    3,075
==========================================================================================
Total params: 6,956,931
Trainable params: 6,956,931
Non-trainable params: 0
Total mult-adds (G): 1.85
==========================================================================================
Input size (MB): 0.20
Forward/backward pass size (MB): 117.90
Params size (MB): 27.83
Estimated Total Size (MB): 145.92
==========================================================================================

We unfortunately do not have a lot of fine-grained control over how this is printed out. However, it is possible to read this into a string, do parsing on our end (e.g., using a separate class) and then print specific portions of the output. This should also help with the out of memory issue that the build is currently facing [ref].

Thoughts?

[sarthakpati]

Updated to print this, which I feel is much cleaner. Thoughts?

[SNIP!]
Device - Current: 0 Count: 1 Name: NVIDIA GeForce RTX 2080 Super with Max-Q Design Availability: True
Model Summary:
        Input size: 0.196664
        Output size: 19.398664
        Parameters size: 36.883972
        Estimated total size: 56.4793
        Total # of operations: 2.446656001 G
[SNIP!]

[codecov]

Codecov Report

Merging #468 (b64ba7e) into master (609371c) will increase coverage by 0.05%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #468      +/-   ##
==========================================
+ Coverage   92.45%   92.51%   +0.05%     
==========================================
  Files         105      105              
  Lines        6298     6348      +50     
==========================================
+ Hits         5823     5873      +50     
  Misses        475      475              

Flag	Coverage Δ
unittests	92.51% <100.00%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
GANDLF/utils/__init__.py	100.00% <ø> (ø)
setup.py	0.00% <ø> (ø)
GANDLF/compute/generic.py	94.44% <100.00%> (+0.32%)	⬆️
GANDLF/parseConfig.py	76.08% <100.00%> (+0.07%)	⬆️
GANDLF/utils/tensor.py	88.34% <100.00%> (+1.01%)	⬆️
testing/test_full.py	98.43% <100.00%> (+0.04%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.