torchlayers is a PyTorch based library providing automatic shape and dimensionality inference of `torch.nn` layers and additional building blocks featured in current SOTA architectures (e.g. Efficient-Net).

Above requires no user intervention (except single call to similarly to the one seen in Keras.

Main functionalities:

  • Shape inference for most of torch.nn module (convolutional, recurrent, transformer, attention and linear layers)

  • Dimensionality inference (e.g. torchlayers.Conv working as torch.nn.Conv1d/2d/3d based on input shape)

  • Shape inference of custom modules (see GitHub README)

  • Additional Keras-like layers (e.g. torchlayers.Reshape or torchlayers.StandardNormalNoise)

  • Additional SOTA layers mostly from ImageNet competitions (e.g. PolyNet, Squeeze-And-Excitation, StochasticDepth)

  • Useful defaults ("same" padding and default kernel_size=3 for Conv, dropout rates etc.)

  • Zero overhead and torchscript support


If you import classes from modules listed belows using torchlayers you will get shape inferrable version, e.g. torchlayers.Conv instead of torchlayers.convolution.Conv.

If you wish to use those without shape inferrence capabilities use fully qualified module name, e.g. torchlayers.convolution.SqueezeExcitation.


Following installation methods are available:


To install latest release:

pip install --user torchlayers

To install nightly version:

pip install --user torchlayers-nightly


Various torchlayers images are available both CPU and GPU-enabled. You can find them at Docker Cloud at szymonmaszke/torchlayers


CPU image is based on ubuntu:18.04 and official release can be pulled with:

docker pull szymonmaszke/torchlayers:18.04

For nightly release:

docker pull szymonmaszke/torchlayers:nightly_18.04

This image is significantly lighter due to lack of GPU support.


All images are based on nvidia/cuda Docker image. Each has corresponding CUDA version tag ( 10.1, 10 and 9.2) CUDNN7 support and base image ( ubuntu:18.04 ).

Following images are available:

  • 10.1-cudnn7-runtime-ubuntu18.04

  • 10.1-runtime-ubuntu18.04

  • 10.0-cudnn7-runtime-ubuntu18.04

  • 10.0-runtime-ubuntu18.04

  • 9.2-cudnn7-runtime-ubuntu18.04

  • 9.2-runtime-ubuntu18.04

Example pull:

docker pull szymonmaszke/torchlayers:10.1-cudnn7-runtime-ubuntu18.04

You can use nightly builds as well, just prefix the tag with nightly_, for example like this:

docker pull szymonmaszke/torchlayers:nightly_10.1-cudnn7-runtime-ubuntu18.04