TorchScan: inspect your PyTorch models

The torchscan package provides tools for analyzing your PyTorch modules and models. Additionally to performance benchmarks, a comprehensive architecture comparison require some insights in the model complexity, its usage of computational and memory resources.

This project is meant for:

• ⚡ exploration: easily assess the influence of your architecture on resource consumption

• 👩‍🔬 research: quickly implement your own ideas to mitigate latency

Supported layers

Here is the list of supported layers for FLOPS, MACs, DMAs and receptive field computation:

Non-linear activations

• torch.nn.ReLU

• torch.nn.ELU

• torch.nn.LeakyReLU

• torch.nn.ReLU6

• torch.nn.Tanh

• torch.nn.Sigmoid

Linear layers

• torch.nn.Identity

• torch.nn.Linear

Convolutions

• torch.nn.Conv1d

• torch.nn.Conv2d

• torch.nn.Conv3d

• torch.nn.ConvTranspose1d

• torch.nn.ConvTranspose2d

• torch.nn.ConvTranspose3d

Pooling

• torch.nn.MaxPool1d

• torch.nn.MaxPool2d

• torch.nn.MaxPool3d

• torch.nn.AvgPool1d

• torch.nn.AvgPool2d

• torch.nn.AvgPool3d

• torch.nn.AdaptiveMaxPool1d

• torch.nn.AdaptiveMaxPool2d

• torch.nn.AdaptiveMaxPool3d

• torch.nn.AdaptiveAvgPool1d

• torch.nn.AdaptiveAvgPool2d

• torch.nn.AdaptiveAvgPool3d

Normalization

• torch.nn.BatchNorm1d

• torch.nn.BatchNorm2d

• torch.nn.BatchNorm3d

Other

• torch.nn.Flatten

• torch.nn.Dropout

Please note that the functional API of PyTorch is not supported.