Pytorch Symbolic: an equivalent of Keras Functional API [Project]

Submitted by gahaalt t3_ypkfwq in MachineLearning

Hello!

I just hit 1.0.0 version of a library I've been developing for the past months as a side project.

Pytorch Symbolic

A library that aims to provide a concise API for neural network creation in PyTorch. The API and the inner workings are similar to Keras / TensorFlow2 Functional API which I always enjoyed using. I decided to go with "Symbolic" in the name instead of "Functional" because I believe it better represents how the library works (also "functional" is kind of taken by torch.nn.functional).

I did my best to prepare a useful documentation, so if you are interested, please check it out! It is filled with examples, best practices, benchmarks, explanations of the inner workings and more.

Example

This example shows how to create a multiple inputs neural network:

from torch import nn
from pytorch_symbolic import Input, SymbolicModel

input1 = Input(shape=(3, 32, 32))
input2 = Input(shape=(3, 32, 32))

output1 = nn.Conv2d(input1.C, 16, 3)(input1)
output2 = nn.Conv2d(input2.C, 16, 3)(input2)
final_output = output1 + output2

model = SymbolicModel(inputs=(input1, input2), outputs=final_output)
model

SymbolicModel(
  (Conv2d_1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1))
  (Conv2d_2): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1))
  (AddOpLayer_1): AddOpLayer()
)

Keras-like summary is available as well:

>>> model.summary()
___________________________________________________________
     Layer          Output shape         Params   Parent   
===========================================================
1    Input_1        (None, 3, 32, 32)    0                 
2    Input_2        (None, 3, 32, 32)    0                 
3    Conv2d_1       (None, 16, 30, 30)   448      1        
4    Conv2d_2       (None, 16, 30, 30)   448      2        
5*   AddOpLayer_1   (None, 16, 30, 30)   0        3,4      
===========================================================
Total params: 896
Trainable params: 896
Non-trainable params: 0
___________________________________________________________

Pytorch Symbolic is actually more powerful than Keras Functional API, because it can be used to define graphs of operations over arbitrary Python objects. So if your torch.nn.Module operates on dictionaries, you can still use it in Pytorch Symbolic. This is not just a gimmick, as it is necessary to be compatible with all the Modules provided by torch.nn and Modules commonly used by the community. You can read more in Advanced Topics section of the documentation.

Installation

Installation is easy with pip and there are no dependencies besides PyTorch (torch>=1.12.0):

pip install pytorch-symbolic

It is a small package, easy to install and uninstall, if you don't like it. :)

There's an introduction in form of Jupyter Notebook, if you prefer it. Go to GitHub to see it or run in Colab.

This library does not compete with the existing:

  • FastAI
  • PyTorch Lightning
  • PyTorch Ignite

All of which are great libraries for data processing and neural network training. Pytorch Symbolic provides API solely for neural network creation and produces models entirely compatible with PyTorch, which means you can train the models created with Pytorch Symbolic using one of the above libraries!

Whether you try it or not, I am excited to hear your feedback on this library. Do you have any suggestions, questions, critique? Please share all your thoughts! :)

Contributions are welcomed too!

151

Comments

You must log in or register to comment.

Bezukhov55 t1_ivjkldm wrote

Cool, maybe will use it in the future) The best thing I like about TF, is that you dont have to specify number of input channels, and here you can use input.C, nice)

15

chatterbox272 t1_ivjvhno wrote

Pytorch does have lazy layers natively, which also don't require known input channels. You just have to do a dummy pass with them after you create a model to init them

13

KingsmanVince t1_ivk562i wrote

Where can I read more about this "lazy layers"?

3

gahaalt OP t1_ivk66uu wrote

Here.

It is just like the standard layer, but retrieves the necessary input shape during the first forward call.

You can use lazy layers in Pytorch Symbolic too.

10

violentdeli8 t1_ivk090t wrote

This is really cool! I am wondering if I can use this as a “compiler” for Neural Architecture Search. During the search phase one has to sample from a search space of architectures and “compile” them into a valid nn.Module. Usually this code is written per search space. Wondering if that part becomes far less tedious with this.

6

gahaalt OP t1_ivk3tvj wrote

Yeah! You have a lot of flexibility to do NAS here. You can create a huge graph of layers and sample a smaller path from it to create a Symbolic Model. One non-standard thing you need to do to pull it off is to modify ._children attribute of Symbolic Data when you want to rewire the connections in this graph.

I might add an example for a simple NAS soon.

4

violentdeli8 t1_ivk53vn wrote

A very simple NAS example using perhaps one of the tabular benchmarks like 201 will be very useful illustration. Thanks!!

2

Zondartul t1_ivju0j9 wrote

Can it do symbolic logic? How does it compare to SymPy?

5

gahaalt OP t1_ivjw8lx wrote

No, it has a different purpose than SymPy. As I understand SymPy is a library mainly for manipulating symbolic mathematical expressions.

Pytorch Symbolic uses symbolic variables to record (capture) the operations and later to replay them on arbitrary data. Under the hood, there's a graph with symbolic variables as nodes and transformations (e.g. layers) as edges.

Pytorch Symbolic can capture and replay arbitrary Python operations, but cannot display them in such neat notation as SymPy does.

6

Mefaso t1_ivn7lv2 wrote

The name is pretty confusing because symbolic already has two meanings in ML:

Symbolic computation, as in sympy, and symbolic AI, as in whatever Gary Marcus likes

This library does neither of the two

3

gahaalt OP t1_ivns08k wrote

Thanks for the opinion. Please look at the article: "What are Symbolic and Imperative APIs in TensorFlow 2.0?" by Josh Gordon linked here. It seems natural to him to describe this API as "Symbolic".

Basically if you google "Symbolic API" it seems to be commonly used to describe this very thing.

Also, a similar nomenclature is used in mxnet.

4

Mefaso t1_ivo2bk2 wrote

Interesting, i didn't know about that. Thanks for sharing the link!

2

metatron7471 t1_iwayuma wrote

The opposite of imperative isn't symbolic but declarative. The name symbolic is very confusing because I also tought it was about (neuro-)symbolic models.

1

SEND_ALL_DOG_PICS t1_ivk1q1f wrote

Why would I use this over torchdynamo? There is a solution built into pytorch for graph capture, which you can then “replay” (in your words) on any data.

5

gahaalt OP t1_ivkmxtm wrote

Thanks for this question!

Pytorch Symbolic is simplifying the definition of the neural network models. It is indeed creating a graph under the hood to do this. In this graph, every edge is an nn.Module.

torchdynamo looks great as a tool for optimizing existing models to perform better on the GPU by removing the CPU overhead entirely. Sometimes the improvement is really impressive.

Yes, torchdynamo does some kind of graph capture as well. It even modifies the byte-code to speed up the execution. But in the end it is a wrapper for an nn.Module that speeds it up. To speed up the model, you have to define it first.

So the two libraries are actually independent. You can use torchdynamo to speed up models created with Pytorch Symbolic. IMO it is a great combination.

12

llun-ved t1_ivlldxy wrote

As a past Keras user, this looks familiar and approachable. In your post, I’d love to see what the “normal” PyTorch construction would be for the simple example.

5

gahaalt OP t1_ivlq9z1 wrote

Thanks a lot for the suggestion! The comparison between symbolic and imperative declarations is indeed interesting. I included it in the documentation, here is the link to the specific section. This is an example of a toy ResNet neural network, still simple but a tad more interesting.

3

gahaalt OP t1_ivlr45z wrote

Let me copy the comparison in case somebody doesn't feel like clicking the link. This might be long, however.

ResNet with the help of Pytorch Symbolic:

from torch import nn
from pytorch_symbolic import Input, SymbolicModel

inputs = Input(shape=(3, 32, 32))
x = nn.Conv2d(inputs.C, 32, 3)(inputs)(nn.ReLU())
x = nn.Conv2d(x.C, 64, 3)(x)(nn.ReLU())
block_1_output = nn.MaxPool2d(3)(x)

x = nn.Conv2d(block_1_output.C, 64, 3, padding=1)(block_1_output)(nn.ReLU())
x = nn.Conv2d(x.C, 64, 3, padding=1)(x)(nn.ReLU())
block_2_output = x + block_1_output

x = nn.Conv2d(block_2_output.C, 64, 3, padding=1)(block_2_output)(nn.ReLU())
x = nn.Conv2d(x.C, 64, 3, padding=1)(x)(nn.ReLU())
block_3_output = x + block_2_output

x = nn.Conv2d(block_3_output.C, 64, 3)(block_3_output)(nn.ReLU())
x = nn.AvgPool2d(kernel_size=x.HW)(x)(nn.Flatten())
x = nn.Linear(x.features, 256)(x)(nn.ReLU())
x = nn.Dropout(0.5)(x)
outputs = nn.Linear(x.features, 10)(x)

model = SymbolicModel(inputs, outputs)

ResNet defined in "standard" PyTorch:

from torch import nn


class ToyResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.relu = nn.ReLU()
        self.block1conv1 = nn.Conv2d(3, 32, 3)
        self.block1conv2 = nn.Conv2d(32, 64, 3)
        self.maxpool = nn.MaxPool2d(3)

        self.block2conv1 = nn.Conv2d(64, 64, 3, padding=1)
        self.block2conv2 = nn.Conv2d(64, 64, 3, padding=1)

        self.block3conv1 = nn.Conv2d(64, 64, 3, padding=1)
        self.block3conv2 = nn.Conv2d(64, 64, 3, padding=1)

        self.conv1 = nn.Conv2d(64, 64, 3)

        kernel_size = 7  # calculated by hand
        self.global_pool = nn.AvgPool2d(kernel_size)
        self.flatten = nn.Flatten()
        self.linear = nn.Linear(64, 256)
        self.dropout = nn.Dropout(0.5)
        self.classifier = nn.Linear(256, 10)

    def forward(self, x):
        x = self.relu(self.block1conv1(x))
        x = self.relu(self.block1conv2(x))
        block_1_output = self.maxpool(x)

        x = self.relu(self.block2conv1(block_1_output))
        x = self.relu(self.block2conv2(x))
        block_2_output = x + block_1_output

        x = self.relu(self.block3conv1(block_2_output))
        x = self.relu(self.block3conv2(x))
        block_3_output = x + block_2_output

        x = self.relu(self.conv1(block_3_output))
        x = self.global_pool(x)
        x = self.flatten(x)
        x = self.relu(self.linear(x))
        x = self.dropout(x)
        return self.classifier(x)


model = ToyResNet()
4

llun-ved t1_ivltagz wrote

This is a great example. Thank you for the relevant link to the docs and this summary.

3

androstudios t1_ivnmdbl wrote

PyTorch has nn.Sequential which allows you to call an entire block as one function

1

gahaalt OP t1_ivnsc23 wrote

Yes. But nn.Sequential won't allow for example for residual connections. You can, however, create a Symbolic Model with residual connections and call the entire model as one function, in your words.

2

bbateman2011 t1_iwiv5jm wrote

As a long-time Keras user, this might be a way to start using PyTorch without having to learn it all at once. That could be a way to "turn the corner" and finally embrace PyTorch in a meaningful work way. What I see is lots of new research stuff is in PyTorch and it's harder for Keras folks to embrace that.

Thanks for sharing this, and all the work that went into ti.

2

VinnyVeritas t1_ivnv1dt wrote

I am not understanding the point of rewriting Keras on top of PyTorch? Is it just for your own fun?

Why not use Keras directly? Keras has better deployment capabilities as a bonus.

0

gahaalt OP t1_ivnveuu wrote

Why not have an API that simplifies model creation in PyTorch? Let's not debate which framework is better. Model creation is but a small brick in the whole framework ecosystem. I am sure there are people who want to stick with PyTorch while creating models conveiniently.

2

VinnyVeritas t1_ivqdjyf wrote

Oh I didn't realize you saw it as a simplification because all it does is remove the input size but then you need to create input placeholders. I thought it was meant to reimplement Keras on PyTorch.

0