Turn TensorFlow functions into mathematical notations and diagrams¶

This is based on some helper classes I started writing, to help my self make less mistakes and understand the code better. It is still work in progress. I want to share the idea to see if it interests others.

TLDR: it is a bunch of helper functions that can act as a direct replacement of tensorflow. It renders mathematical formulas and output tensor dimensions. I'm looking for feedback.

Example¶

There's a bunch of things I've been working on. In order to show them I've used the initial model creation code of this tensorflow implementation of capsule nets. Starting from Input Images upto Digit Capsules.

Import the library. I'm trying to make it work as a direct replacement.

import tf_for_beginners as tf

All the operations and variables are assigned to a model. This takes care of variable names, so that you don't have to explicitly specify name=.

model = tf.Model()

Create a place holder just like you do with TensorFlow, and assign it to model. The library will set the tensorflow name, and output a nice 3D diagram showing the dimensions of the placeholder.

model.X = tf.placeholder(shape=[None, 28, 28, 1], dtype=tf.float32)

The dimensions of the tensor is visualized in 3D. This is more convenient than having to write an extra line to output the shape of $X$. I'm not sure if this is the best way to visualize; I have discussed some other ideas later.

Also notice that we didn't pass the variable name as an argument; it was extracted from the assigned property name. The cost of this is that we have to use a model to store all variables and operations.

caps1_n_maps = 32
caps1_n_caps = caps1_n_maps * 6 * 6  # 1152 primary capsules
caps1_n_dims = 8

conv1_params = {
    "filters": 256,
    "kernel_size": 9,
    "strides": 1,
    "padding": "valid",
    "activation": tf.nn.relu,
}

conv2_params = {
    "filters": caps1_n_maps * caps1_n_dims, # 256 convolutional filters
    "kernel_size": 9,
    "strides": 2,
    "padding": "valid",
    "activation": tf.nn.relu
}

You can add conv2d layers similarly.

model.conv1 = tf.layers.conv2d(model.X, **conv1_params)
model.conv2 = tf.layers.conv2d(model.conv1, **conv2_params)

Above diagram shows that there are 6x6 feature maps of 256 channels after the second convolution. I'm planning to indicate the strides, kernel size and padding on the 2-D diagram.

model.caps1_raw = tf.reshape(model.conv2, [-1, caps1_n_caps, caps1_n_dims])

Here's a function definition. You can add a naming scope, so that it will not render internal operations.

def squash(s, axis=-1, epsilon=tf.constant(1e-7), name=None):
    model.epsilon = epsilon
    if name is None:
        name = 'squash'
    model.add_scope(name)
    model.squared_norm = tf.reduce_sum(tf.square(s), axis=axis,
                                 keep_dims=True)
    model.safe_norm = tf.sqrt(tf.add(model.squared_norm, epsilon))
    model.squash_factor = tf.divide(model.squared_norm, tf.add(1., model.squared_norm))
    model.unit_vector = tf.divide(s, model.safe_norm)
    model.squashed = tf.multiply(model.squash_factor, model.unit_vector)

    model.remove_scope() # Need to implement with mode.scope ....

    return model.squashed

Call to the function will render the full formula and output tensor dimensions.

model.caps1_output = squash(model.caps1_raw, name="caps1_squash")

This helps you verify the formulas when coding. It help readers quickly understand whats going on, without having to look back at the function definition.

Notice that, tries to convert the variable names into mathematical symbols; e.g. epsilon -> $\epsilon$

Why?¶

This was developed to make machine learning code easier to understand, at the time of writing as well as for readers. I have only tried this on small networks (mostly based on tutorials). So far I've found it helpful.

I feel it's helpful for beginners than for experts, who get confused on the detailed workings of the functions. This is espcially true for me because I only work on tensorflow on and off, and I forget all the details.

It is also useful when writing notebooks intended for others to read (like tutorials).

What's next?¶

The above example shows most of the things I've implemented so far. It's just a couple of hunderd lines of python and javascript I've written on free time. I am planning to cleanup the code and host it on github. Right now the code is messy, and needs to be fixed almost everytime I use it.

Cover more tensorflow operations¶

Current implementaion is only for a few common tensorflow functions.

Also, the library doesn't do operator overloading which is really important. That is, we should be able to write X / Y instead of tf.divide(X, Y)

Diagrams to how operations work¶

We need to include more details in diagrams to explain complex operation. For instance, we can indicate padding, strides and kernal size in the diagram for convolution operation. We can do similar diagrams for tensor transformations.

Also, the current 2-D and 3-D tensor diagrams might not be the best in all situations. It sometimes makes more sense not to show individual cells, butto show the sizes of the dimensions.

Why I decided to show individual cells is with the intention of adding details like highlighted cells, kernels, etc on the same diagram.

matrix

It's sometimes important to show higher dimensianality, especially for operations like reshaping and tiling.

tiling

And, on some occasions, it may be just enough to show the shape of the tensor in text; e.g. ([?, 10, 50, 20, 10]).

Function Wrappers¶

For example, we can wrap the function squash.

wrap(tensor(['N', 'H', 'W'])
def squash(s, axis=-1, epsilon=Noob(tf.constant, 1e-7), name=None):
    ...

And it can display the full formula for the function and the output tensor dimensions. This will help us test functions before calling them, and help readers get an overview of the function without reading the internals.