Generating a *simple* TensorFlow graph illustration - tensorflow

Question

I'm working on my first deep learning model using TensorFlow in a Jupyter notebook, and I would like to generate simplified graphs which illustrate the various layers of the network. Specifically, graphs such as those pictured in this answer:
This is very simple and clean and I can understand what's going on. This is more important than capturing 100% of the details. Contrast with the graph generated by TensorBoard which is a complete fustercluck:
How can I take a tf.Graph object and automatically generate a graph similar to the one above? Bonus points if it can be displayed in the Jupyter Notebook, too.

Answer 1

In short - you cannot. TF is a low-level library, which has no concept of "high level operations", it has ops, and this is the only thing it can visualise in a way you are thinking about. In particular, from math perspective there are no "neurons" in your graph, there are just tensors being multiplied by each other, this additional "semantics" is there only to make it easier for humans to talk about this, but it is not really encoded in your graph.
What you can do is to group nodes by yourself by specifing variable_scope for sections of your graph, then, after displaying in TB they will be displayed as a single node. It will not give you this "per-neuron-like" flavour of visualisation but at least it will hide many details. Creating a nice, visually appealing visualisations of neural nets is an "art" on its own rights, and a hard task to do in general.

Answer 2

Here's a snippet of code that we use in our PipelineAI notebooks to display our TensorFlow graphs inline within our Jupyter notebooks:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import re
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2
def convert_graph_to_dot(input_graph, output_dot, is_input_graph_binary):
graph = graph_pb2.GraphDef()
with open(input_graph, "rb") as fh:
if is_input_graph_binary:
graph.ParseFromString(fh.read())
else:
text_format.Merge(fh.read(), graph)
with open(output_dot, "wt") as fh:
print("digraph graphname {", file=fh)
for node in graph.node:
output_name = node.name
print(" \"" + output_name + "\" [label=\"" + node.op + "\"];", file=fh)
for input_full_name in node.input:
parts = input_full_name.split(":")
input_name = re.sub(r"^\^", "", parts[0])
print(" \"" + input_name + "\" -> \"" + output_name + "\";", file=fh)
print("}", file=fh)
print("Created dot file '%s' for graph '%s'." % (output_dot, input_graph))
input_graph='/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb'
output_dot='/root/notebooks/unoptimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)
Using graphviz, you can convert the .dot to .png using a %%bash magic within your notebook cell:
%%bash
dot -T png /root/notebooks/unoptimized_cpu.dot \
-o /root/notebooks/unoptimized_cpu.png > /tmp/a.out
and finally, display the graph in your notebook:
from IPython.display import Image
Image('/root/notebooks/unoptimized_cpu.png', width=1024, height=768)
here's an example of a simple Linear Regression model implemented in TensorFlow:
Here's the optimized version used to deploy and serve the TensorFlow Model in production (also rendered using the above code snippets):
More examples and details of these types of optimizations at http://pipeline.ai