Build PyTorch Models Easily Using torchlayers

torchlayers aims to do what Keras did for TensorFlow, providing a higher-level model-building API and some handy defaults and add-ons useful for crafting PyTorch neural networks.



PyTorch continues to enjoy an impressive level of general interest, judging by online searches and, more importantly, continues to increase its rate of adoption. PyTorch is thought of as powerful and flexible, a pair of characteristics which are welcome by researchers. However, PyTorch has, in the past, been susceptible to criticism by practitioners for its lack of a simplified higher-level API, such as TensorFlow's Keras. This situation has recently changed.

Figure

 

torchlayers aims to do for PyTorch what Keras has done for TensorFlow. Concisely defined via the project's developers:

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

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

Aside from the shape and dimensionality inference mentioned above, torchlayers includes additional Keras-like layers such as torchlayers.Reshape (reshapes input tenors while preserving batch dimensions), includes SOTA layers previously seen in ImageNet competitions (PolyNet, for example), and provides some useful defaults such as convolution kernel size (the default in torchlayers is 3).

Installation is simple using pip:

pip install --user torchlayers

 

Additional installation information (re: Docker images and GPUs) can be found here. Full torchlayers documentation can be found here.

The torchlayers GitHub page provides a few examples to show off some of its features. I like the Simple image and text classifier in one! example, the code for which I have replicated below. This example demonstrates:

  • torch.nn and torchlayers layer intermixing
  • shape and dimensionality inference (Conv and Linear input and BatchNorm)
  • default Conv kernel size
  • Conv padding defaults to "same"
  • use of torchlayers pooling layer (GlobalMaxPool, similar to Keras)
import torch
import torchlayers as tl

# torch.nn and torchlayers can be mixed easily
model = torch.nn.Sequential(
    tl.Conv(64),                   # specify ONLY out_channels
    torch.nn.ReLU(),               # use torch.nn wherever you wish
    tl.BatchNorm(),                # BatchNormNd inferred from input
    tl.Conv(128),                  # Default kernel_size equal to 3
    tl.ReLU(),
    tl.Conv(256, kernel_size=11),  # "same" padding as default
    tl.GlobalMaxPool(),            # Known from Keras
    tl.Linear(10),                 # Output for 10 classes
)

 

A defined network can then be built using torchlayers.build while specifying input shape (input shapes shown below for both image and text classification, for the model defined above):

# Image...
mnist_model = tl.build(model, torch.randn(1, 3, 28, 28))

# ...or text
# [batch, embedding, timesteps], first dimension > 1 for BatchNorm1d to work
text_model = tl.build(model, torch.randn(2, 300, 1))

 

build obviously works like in Keras, performing the equivalent of compiling the model to PyTorch primitives; it allows for some additional functionality via a post_build function (such as weight initialization, shown below), which you can find more about here.

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

 

torchlayers provides some helpful features for Keras-like model building using PyTorch, and fills an obvious hole. Time will tell how the project develops and catches on long-term, but it's certainly off to a promising start.

 
Related: