A very lightweight and minimalistic output shape examiner of layers and models.
Project description
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
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
output-shape-0.0.1.tar.gz
(3.1 kB
view hashes)
Built Distribution
Close
Hashes for output_shape-0.0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 59977aec503b0655d6ddbd04d99f8f70799c8ae5f579070eb93819131996cb32 |
|
MD5 | 543bc46451e186f74c640711085b3842 |
|
BLAKE2b-256 | 5c23badff3b1d8313c252aa6b8721121c7cb442113b47fb816f09f9fce33a57b |