Visualizing TensorFlow Graphs in Jupyter Notebooks

Prerequisites: This article assumes you are familiar with the basics of Python, TensorFlow, and Jupyter notebooks.
We won't use any of the advanced TensorFlow features, as our goal is just to visualize the computation graphs.

TensorFlow operations form a computation graph. And while for small examples
you might be able to look at the code and immediately see what is going on,
larger computation graphs might not be so obvious. Visualizing the graph can help
both in diagnosing issues with the computation itself, but also in understanding
how certain operations in TensorFlow work and how are things put together.

We'll take a look at a few different ways of visualizing TensorFlow graphs, and most importantly,
show how to do it in a very simple and time-efficient way. It shouldn't take more than one or two lines
of code to draw a graph we have already defined. Now onto the specifics, we'll take a look at the following
visualization techniques:

Exploring the textual graph definition

Building a GraphViz DOTgraph from that directly in the Jupyter Notebook

Visualizing the same graph in a locally running instance of TensorBoard

Using a self contained snippet that uses a cloud deployed publically available TensorBoard instance to render the graph inline in a Jupyter Notebook.

First, let us create a simple TensorFlow graph. Regular operations such as creating a placeholder with tf.placeholder will create a node in the so called default graph. We can access it via tf.get_default_graph(), but we can also change it temporarily. In our example below, we'll create a new instance of the tf.Graph object and create a simple operation adding two variables

$$c = a + b$$

Note that we're giving explicit names to both of the placeholder variables.

The variable g now contains a definition of the computation graph for the operation $c = a + b$.
We can use the g.as_graph_def() method to get a textual representation of the graph for our expression.
While the main use of this is for serialization and later deserialization via tf.import_graph_def, we'll use it to create a GraphViz DOTgraph.

Let us take a look at the GraphDef for our simple
expression. First, we'll inspect the names of all of the nodes in the graph.

In [35]:

[node.namefornodeing.as_graph_def().node]

Out[35]:

['a', 'b', 'add']

As expected, there are three nodes in the Graph. One for each of our variables, and one for the addition
opeartion. The placeholder variable nodes have a name since we explicitely named them when calling tf.placeholder. If we omit the name keyword argument, TensorFlow will simply generate a name on its own, as it did with the add operation.

Next, we can take a look at the edges in the graph. Each GraphDef node has an input field which specifies
of the nodes where it has edges. Let's take a look:

In [36]:

g.as_graph_def().node[2].input

Out[36]:

['a', 'b']

As we can see, there are two edges, one to each variable. We can feed this directly into GraphViz.

GraphViz is a fairly popular library for drawing graphs, trees and other graph-shaped data structures. We'll use the Python GraphViz package which provides a nice
clean interface. We can install it directly inside a Jupyter notebook via !pip install graphviz.

The graph definition itself will be rather simple, and we'll take inspiration from a similar piece of code in
TensorFlow itself (in graph_to_dot.py) which generates a DOTgraph file format for a given GraphDef. Unfortunately it is only available as a command line script, and as such we can't call it directly from our code. This is why we'll be implementing it ourselves, but don't worry, it will only
be a few lines of code.

In [56]:

fromgraphvizimportDigraphdot=Digraph()forning.as_graph_def().node:# Each node has a name and a label. The name identifies the node# while the label is what will be displayed in the graph.# We're using the name as a label for simplicity.dot.node(n.name,label=n.name)foriinn.input:# Edges are determined by the names of the nodesdot.edge(i,n.name)# Jupyter can automatically display the DOT graph,# which allows us to just return it as a value.dot

Out[56]:

%3aaaddadda->addbbb->add

Now let's wrap this in a function and try using it on a more complicated expression.

We'll build another graph calculating the area of a circle with the formula $\pi * r^2$.
As we can see TensorFlow does what we would actually expect and links the same placeholder
to two multiplication operations.

While GraphViz might be nice for visualizing small graphs, neural networks can grow to quite a large size.
TensorBoard allows us to easily group parts of our equations into scopes, which will then be visually
separated in the resulting graph. But before doing this, let's just try visualizing our previous graph
with TensorBoard.

All we need to do is save it using the tf.summary.FileWriter, which takes a directory and a graph, and serializes
the graph in a format that TensorBoard can read. The directory can be anything you'd like, just make sure you point to the same directory using the tensorboard --logdir=DIR command (DIR being the directory you specified to the FileWriter).

In [65]:

# We write the graph out to the `logs` directorytf.summary.FileWriter("logs",g).close()

Next, open up a console and navigate to the same directory from which you executed the FileWriter command, and run tensorboard --logdir=logs. This will launch an instance of TensorBoard which you can access at http://localhost:6006. Then navigate to the Graphs section and you should see a graph that looks like the following image. Note that you can also click on the nodes in the graph to inspect them further.

Now this is all nice and interactive, but we can already see some things which make it harder to read. For example, when we type $\pi * r^2$ we generally don't think of the $r^2$ as a multiplication operation (even though we implement it as such), we think of it as a square operation. This becomes more visible when the graph contains a lot more
operations.

Luckily, TensorFlow allows us to bundle operations together into a single unit called scope. But first, lets take a look at a more complicated example without using scopes. We'll create a very simple feed forward neural network with three layers (with respective weights $W_1, W_2, W_3$ and biases $b_1, b_2, b_3$).

Looking at the result in TensorBoard, the result is pretty much what we would expect. The only problem is, TensorBoard displays it as a single expression. It isn't immediately apparent that we meant to think about our code in terms of layers.

We can improve this by using the above-mentioned tf.name_scope
function. Let us rewrite our feedforward network code to separate each layer into its own scope.

We'll use the modified snippet from the DeepDream notebook
taken from this StackOverflow answer. It basically takes the tf.GraphDef,
sends it over to the cloud, and embeds an <iframe> with the resulting visualization right in the Jupyter notebook.

Here's the snippet in its whole. All you need to do is call show_graph() and it will handle everything, as shown in the
example below on our previous graph g. The obvious advantage of this approach is that you don't need to run TensorBoard
to visualize the data, but you also need internet access.

Thats it for this article! Hopefully this article showed you a few tricks that can help you solve TensorFlow problems more
effectively. Here are a few links to related articles and references that further describe TensorBoard.