PyTorch model summaries

Improvements:

modernizes and simplifies the code
adds input shape tracking
supports arbitrary input objects to the model (through the extract_input_tensor argument to summary())
adds Conv2d complexity
adds start, jump, receptive field tracking (based on pytorch-receptive-field and this medium article)
adds receptive-field computation based on gradients (see this notebook)

Demo

pytest --nbmake demo.ipynb

For complex models, this will likely be accurate where the analytical computation is not: for ex. when two modules’ outputs are combined

grad_receptive_field=True modifies the model in multiple ways (but makes a copy before doing so)
- initialization of Conv2d
- replaces MaxPool2d with AvgPool2d
- turns off Dropout and BatchNorm2d
Treats any class ending in Conv2d as a Conv2d, and similarly for BatchNorm2d, MaxPool2d, Dropout. This is for handling custom module classes that don’t directly derive from the nn. classes
Requires a tensor output from a layer to compute RF for that layer

Computed as num_input_filters * num_output_filters * H * W

Handle AdaptiveAvgPool2d: https://stackoverflow.com/a/63603993/14044156
See if we can generate graphviz dot, either by using FX (https://github.com/pytorch/examples/blob/master/fx/module_tracer.py) or through hooks itself?