torchlayers package¶

torchlayers.build(module, *args, **kwargs)[source]¶

Build PyTorch layer or module by providing example input.

This method should be used always after creating module using torchlayers and shape inference especially.

Works similarly to build functionality provided by keras.

Provided module will be “compiled” to PyTorch primitives to remove any overhead.

torchlayers also supports post_build function to perform some action after shape was inferred (weight initialization example below):

import torch
import torchlayers as tl

class _MyModuleImpl(torch.nn.Linear):
    def post_build(self):
        # You can do anything here really
        torch.nn.init.eye_(self.weights)

MyModule = tl.infer(_MyModuleImpl)

post_build should have no arguments other than self so all necessary data should be saved in module beforehand.

Parameters

module (torch.nn.Module) – Instance of module to build
*args – Arguments required by module’s forward
**kwargs – Keyword arguments required by module’s forward

torchlayers.infer(module_class, index: str = 1)[source]¶

Allows custom user modules to infer input shape.

Input shape should be the first argument after self.

Usually used as class decorator, e.g.:

import torch
import torchlayers as tl

class _StrangeLinearImpl(torch.nn.Linear):
    def __init__(self, in_features, out_features, bias: bool = True):
        super().__init__(in_features, out_features, bias)
        self.params = torch.nn.Parameter(torch.randn(out_features))

    def forward(self, inputs):
        super().forward(inputs) + self.params

# Now you can use shape inference of in_features
StrangeLinear = tl.infer(_StrangeLinearImpl)

# in_features can be inferred
layer = StrangeLinear(out_features=64)

Parameters

module_class (torch.nn.Module) – Class of module to be updated with shape inference capabilities.
index (int, optional) – Index into tensor.shape input which should be inferred, e.g. tensor.shape[1]. Default: 1 (0 being batch dimension)

class torchlayers.Lambda(function: Callable)[source]¶

Use any function as torch.nn.Module

Simple proxy which allows you to use your own custom in torch.nn.Sequential and other requiring torch.nn.Module as input:

import torch
import torchlayers as tl

model = torch.nn.Sequential(tl.Lambda(lambda tensor: tensor ** 2))
model(torch.randn(64 , 20))

Parameters: function (Callable) – Any user specified function
Returns: Anything function returns
Return type: Any

forward(*args, **kwargs) → Any[source]¶

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchlayers.Reshape(*shapes: int)[source]¶

Reshape tensor excluding batch dimension

Reshapes input torch.Tensor features while preserving batch dimension. Standard torch.reshape values (e.g. -1) are supported, e.g.:

import torch
import torchlayers as tl

layer = tl.Reshape(20, -1)
layer(torch.randn(64, 80)) # shape (64, 20, 4)

All tensors must have the same shape (except in the concatenating dimension). If possible, no copy of tensor will be performed.

Parameters: shapes (*int) – Variable length list of shapes used in view function
Returns: Concatenated tensor
Return type: torch.Tensor

forward(tensor)[source]¶

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchlayers.Concatenate(dim: int)[source]¶

Concatenate list of tensors.

Mainly useful in torch.nn.Sequential when previous layer returns multiple tensors, e.g.:

import torch
import torchlayers as tl

class Foo(torch.nn.Module):
    # Return same tensor three times
    # You could explicitly return a list or tuple as well
    def forward(tensor):
        return tensor, tensor, tensor


model = torch.nn.Sequential(Foo(), tl.Concatenate())
model(torch.randn(64 , 20))

All tensors must have the same shape (except in the concatenating dimension).

Parameters: dim (int) – Dimension along which tensors will be concatenated
Returns: Concatenated tensor along specified dim.
Return type: torch.Tensor

forward(inputs)[source]¶

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

torchlayers package¶

Submodules¶