output-shape 0.0.1

Description:

outputshape 0.0.1

output-shape

A very lightweight and minimalistic output shape examiner of layers and models.
** Currently working for PyTorch models only. **
Installation
pip install output-shape

Usage
You need to decorate the forward method of your model with the decorator and add a debug flag to the init of your model.
import torch
import output_shape

class Model(torch.nn.Module):
def __init__(self, debug=False):
self.debug = debug
...

@output_shape
def forward(self, x):
...

model = Model(debug=True)(torch.randn(2, 1, 128, 128))

Input torch.Size([2, 1, 128, 128])
Conv2d torch.Size([2, 768, 8, 8])
PatchEmbed torch.Size([2, 64, 768])
LayerNorm torch.Size([2, 13, 768])
Linear torch.Size([2, 13, 2304])
Linear torch.Size([2, 13, 768])
Dropout torch.Size([2, 13, 768])
Attention torch.Size([2, 13, 768])
PreNorm torch.Size([2, 13, 768])
LayerNorm torch.Size([2, 13, 768])
Linear torch.Size([2, 13, 3072])
GELU torch.Size([2, 13, 3072])
Dropout torch.Size([2, 13, 3072])
Linear torch.Size([2, 13, 768])
Dropout torch.Size([2, 13, 768])
FeedForward torch.Size([2, 13, 768])
PreNorm torch.Size([2, 13, 768])
Transformer torch.Size([2, 13, 768])
LayerNorm torch.Size([2, 13, 768])
Linear torch.Size([2, 12, 512])
LayerNorm torch.Size([2, 8, 8, 512])
CyclicShift torch.Size([2, 8, 8, 512])
Linear