torchlayers

torchlayers

torchlayers is a PyTorch-based library that provides automatic shape and dimensionality inference for torch.nn layers. It allows you to define models by specifying only the output dimensions (e.g., out_channels or out_features), while the library automatically infers the input dimensions from an example input during instantiation.

Key Features:

• Shape Inference: Works for most torch.nn modules including convolutional, recurrent, transformer, attention, and linear layers.
• Dimensionality Inference: Automatically selects the correct dimensionality (e.g., tl.Conv becomes Conv1d, Conv2d, or Conv3d based on the input shape).
• Keras-like API: Provides additional building blocks like tl.Reshape or tl.StandardNormalNoise and uses sensible defaults (e.g., "same" padding and kernel_size=3 for Conv).
• SOTA Layers: Includes layers from modern architectures like PolyNet, Squeeze-And-Excitation, and StochasticDepth.
• Compatibility: Works seamlessly with standard torch.nn.Module and torch.nn.Sequential, and supports torchscript with zero overhead.

Requirements:

• Python >= 3.7
• PyTorch >= 1.3.0

The primary purpose of torchlayers is to provide automatic shape and dimensionality inference for torch.nn layers, similar to the Keras API. This allows you to build models without manually calculating input/output dimensions for every layer.

Key Capabilities:

• Shape Inference: Automatically handles most torch.nn modules, including convolutional, recurrent, transformer, attention, and linear layers.
• Dimensionality Inference: Layers like torchlayers.Conv automatically determine whether to act as torch.nn.Conv1d, Conv2d, or Conv3d based on the provided input shape.
• Custom Module Inference: You can define your own modules with shape inference capabilities.
• Keras-like Layers: Provides additional building blocks such as torchlayers.Reshape or torchlayers.StandardNormalNoise.
• SOTA Layers: Includes layers from state-of-the-art architectures (e.g., PolyNet, Squeeze-And-Excitation, StochasticDepth).
• Zero Overhead: The inference process does not add runtime overhead and supports torchscript.

Usage Note: To use the shape-inferrable versions of layers, import them directly from the torchlayers namespace (e.g., torchlayers.Conv). If you need to use a specific module without shape inference, use its fully qualified name (e.g., torchlayers.convolution.SqueezeExcitation).

To instantiate a model where dimensions are automatically inferred, define your architecture using torchlayers modules (or standard torch.nn modules) and then call torchlayers.build(model, example_input).

When using torchlayers wrappers like tl.Conv2d or tl.Linear, you only need to provide the output dimensions. The input dimensions will be determined by the shape of the example_input provided to tl.build.

Example: Basic Classifier

import torch
import torchlayers as tl

class Classifier(tl.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = tl.Conv2d(64, kernel_size=6)
        self.conv2 = tl.Conv2d(128, kernel_size=3)
        self.conv3 = tl.Conv2d(256, kernel_size=3, padding=1)
        self.pooling = tl.GlobalMaxPool()
        self.dense = tl.Linear(10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.relu(self.conv3(x))
        return self.dense(self.pooling(x))

# Pass model and any example inputs afterwards
clf = tl.build(Classifier(), torch.randn(1, 3, 32, 32))

Various Docker images are available for both CPU and GPU environments.

CPU Images

CPU images are based on ubuntu:18.04 and are lighter because they lack GPU support.

To pull the official CPU release:

docker pull szymonmaszke/torchlayers:18.04

To pull the nightly CPU release:

docker pull szymonmaszke/torchlayers:nightly_18.04

GPU Images

GPU images are based on nvidia/cuda and support various CUDA versions (10.1, 10, and 9.2) and CUDNN7.

Available tags include:

• 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 for a specific GPU runtime:

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

To use nightly GPU builds, prefix the tag with nightly_ (e.g., nightly_10.1-cudnn7-runtime-ubuntu18.04).

You can combine standard torch.nn modules with torchlayers modules in a single Sequential container or custom Module. torchlayers will handle the inference for its own components, while standard PyTorch modules will behave as usual.

One powerful feature is that a single model definition can be used for different tasks (e.g., Image vs. Text) simply by passing different example input shapes to tl.build. torchlayers will automatically switch the dimensionality of layers like Conv and BatchNorm to match the input.

Example: Multi-purpose Model

import torch
import torchlayers as tl

# Define a model using both torch.nn and torchlayers
model = torch.nn.Sequential(
    tl.Conv(64),  # specify ONLY out_channels
    torch.nn.ReLU(),
    tl.BatchNorm(),  # BatchNormNd inferred from input
    tl.Conv(128),
    tl.ReLU(),
    tl.Conv(256, kernel_size=11),
    tl.GlobalMaxPool(),
    tl.Linear(10),
)

# Build for MNIST (Image classification)
# Input shape: [batch, channels, height, width]
mnist_model = tl.build(model, torch.randn(1, 3, 28, 28))

# Build the SAME model for Text (Sequence classification)
# Input shape: [batch, embedding, timesteps]
# Note: first dimension > 1 for BatchNorm1d to work
text_model = tl.build(model, torch.randn(2, 300, 1))

You can make any custom torch.nn.Module shape-inferable by using torchlayers.infer().

By default, tl.infer uses inputs.shape[1] as the value for the dimension to be inferred. If you need to infer from a different dimension (e.g., the third dimension), you can specify the index argument.

Implementation Pattern: It is recommended to define your base implementation with a prefix like _ and a postfix like Impl to distinguish it from the shape-inferable version.

Example: Custom Linear Layer

import torch
import torchlayers as tl

# 1. Define the base implementation
class _MyLinearImpl(torch.nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))
        self.bias = torch.nn.Parameter(torch.randn(out_features))

    def forward(self, inputs):
        return torch.nn.functional.linear(inputs, self.weight, self.bias)

# 2. Create the shape-inferable version
MyLinear = tl.infer(_MyLinearImpl)

# 3. Build and use
layer = tl.build(MyLinear(out_features=32), torch.randn(1, 64))
output = layer(torch.randn(1, 64))