Simple way to visualize a TensorFlow graph in Jupyter? - tensorflow

Question

The official way to visualize a TensorFlow graph is with TensorBoard, but sometimes I just want a quick look at the graph when I'm working in Jupyter.
Is there a quick solution, ideally based on TensorFlow tools, or standard SciPy packages (like matplotlib), but if necessary based on 3rd party libraries?

Answer 1

Here's a recipe I copied from one of Alex Mordvintsev deep dream notebook at some point
from IPython.display import clear_output, Image, display, HTML
import numpy as np
def strip_consts(graph_def, max_const_size=32):
"""Strip large constant values from graph_def."""
strip_def = tf.GraphDef()
for n0 in graph_def.node:
n = strip_def.node.add()
n.MergeFrom(n0)
if n.op == 'Const':
tensor = n.attr['value'].tensor
size = len(tensor.tensor_content)
if size > max_const_size:
tensor.tensor_content = "<stripped %d bytes>"%size
return strip_def
def show_graph(graph_def, max_const_size=32):
"""Visualize TensorFlow graph."""
if hasattr(graph_def, 'as_graph_def'):
graph_def = graph_def.as_graph_def()
strip_def = strip_consts(graph_def, max_const_size=max_const_size)
code = """
<script>
function load() {{
document.getElementById("{id}").pbtxt = {data};
}}
</script>
<link rel="import" href="https://tensorboard.appspot.com/tf-graph-basic.build.html" onload=load()>
<div style="height:600px">
<tf-graph-basic id="{id}"></tf-graph-basic>
</div>
""".format(data=repr(str(strip_def)), id='graph'+str(np.random.rand()))
iframe = """
<iframe seamless style="width:1200px;height:620px;border:0" srcdoc="{}"></iframe>
""".format(code.replace('"', '"'))
display(HTML(iframe))
Then to visualize current graph
show_graph(tf.get_default_graph().as_graph_def())
If your graph is saved as pbtxt, you could do
gdef = tf.GraphDef()
from google.protobuf import text_format
text_format.Merge(open("tf_persistent.pbtxt").read(), gdef)
show_graph(gdef)
You'll see something like this

Answer 2

TensorFlow 2.0 now supportsTensorBoardinJupytervia magic commands (e.g %tensorboard --logdir logs/train). Here's a link to tutorials and examples.
[EDITS 1, 2]
As #MiniQuark mentioned in a comment, we need to load the extension first(%load_ext tensorboard.notebook).
Below are usage examples for using graph mode, #tf.function and tf.keras (in tensorflow==2.0.0-alpha0):
1. Example using graph mode in TF2 (via tf.compat.v1.disable_eager_execution())
%load_ext tensorboard.notebook
import tensorflow as tf
tf.compat.v1.disable_eager_execution()
from tensorflow.python.ops.array_ops import placeholder
from tensorflow.python.training.gradient_descent import GradientDescentOptimizer
from tensorflow.python.summary.writer.writer import FileWriter
with tf.name_scope('inputs'):
x = placeholder(tf.float32, shape=[None, 2], name='x')
y = placeholder(tf.int32, shape=[None], name='y')
with tf.name_scope('logits'):
layer = tf.keras.layers.Dense(units=2)
logits = layer(x)
with tf.name_scope('loss'):
xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=logits)
loss_op = tf.reduce_mean(xentropy)
with tf.name_scope('optimizer'):
optimizer = GradientDescentOptimizer(0.01)
train_op = optimizer.minimize(loss_op)
FileWriter('logs/train', graph=train_op.graph).close()
%tensorboard --logdir logs/train
2. Same example as above but now using #tf.function decorator for forward-backward passes and without disabling eager execution:
%load_ext tensorboard.notebook
import tensorflow as tf
import numpy as np
logdir = 'logs/'
writer = tf.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True, profiler=True)
#tf.function
def forward_and_backward(x, y, w, b, lr=tf.constant(0.01)):
with tf.name_scope('logits'):
logits = tf.matmul(x, w) + b
with tf.name_scope('loss'):
loss_fn = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=y, logits=logits)
reduced = tf.reduce_sum(loss_fn)
with tf.name_scope('optimizer'):
grads = tf.gradients(reduced, [w, b])
_ = [x.assign(x - g*lr) for g, x in zip(grads, [w, b])]
return reduced
# inputs
x = tf.convert_to_tensor(np.ones([1, 2]), dtype=tf.float32)
y = tf.convert_to_tensor(np.array([1]))
# params
w = tf.Variable(tf.random.normal([2, 2]), dtype=tf.float32)
b = tf.Variable(tf.zeros([1, 2]), dtype=tf.float32)
loss_val = forward_and_backward(x, y, w, b)
with writer.as_default():
tf.summary.trace_export(
name='NN',
step=0,
profiler_outdir=logdir)
%tensorboard --logdir logs/
3. Using tf.keras API:
%load_ext tensorboard.notebook
import tensorflow as tf
import numpy as np
x_train = [np.ones((1, 2))]
y_train = [np.ones(1)]
model = tf.keras.models.Sequential([tf.keras.layers.Dense(2, input_shape=(2, ))])
model.compile(
optimizer='sgd',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
logdir = "logs/"
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)
model.fit(x_train,
y_train,
batch_size=1,
epochs=1,
callbacks=[tensorboard_callback])
%tensorboard --logdir logs/
These examples will produce something like this below the cell:

Answer 3

I wrote a Jupyter extension for tensorboard integration. It can:
Start tensorboard just by clicking a button in Jupyter
Manage multiple tensorboard instances.
Seamless integration with Jupyter interface.
Github: https://github.com/lspvic/jupyter_tensorboard

Answer 4

I wrote a simple helper which starts a tensorboard from the jupyter notebook. Just add this function somewhere at the top of your notebook
def TB(cleanup=False):
import webbrowser
webbrowser.open('http://127.0.1.1:6006')
!tensorboard --logdir="logs"
if cleanup:
!rm -R logs/
And then run it TB() whenever you generated your summaries. Instead of opening a graph in the same jupyter window, it:
starts a tensorboard
opens a new tab with tensorboard
navigate you to this tab
After you are done with exploration, just click the tab, and stop interrupt the kernel. If you want to cleanup your log directory, after the run, just run TB(1)

Answer 5

A Tensorboard / iframes free version of this visualization that admittedly gets cluttered quickly can
import pydot
from itertools import chain
def tf_graph_to_dot(in_graph):
dot = pydot.Dot()
dot.set('rankdir', 'LR')
dot.set('concentrate', True)
dot.set_node_defaults(shape='record')
all_ops = in_graph.get_operations()
all_tens_dict = {k: i for i,k in enumerate(set(chain(*[c_op.outputs for c_op in all_ops])))}
for c_node in all_tens_dict.keys():
node = pydot.Node(c_node.name)#, label=label)
dot.add_node(node)
for c_op in all_ops:
for c_output in c_op.outputs:
for c_input in c_op.inputs:
dot.add_edge(pydot.Edge(c_input.name, c_output.name))
return dot
which can then be followed by
from IPython.display import SVG
# Define model
tf_graph_to_dot(graph).write_svg('simple_tf.svg')
SVG('simple_tf.svg')
to render the graph as records in a static SVG file

Answer 6

Code
def tb(logdir="logs", port=6006, open_tab=True, sleep=2):
import subprocess
proc = subprocess.Popen(
"tensorboard --logdir={0} --port={1}".format(logdir, port), shell=True)
if open_tab:
import time
time.sleep(sleep)
import webbrowser
webbrowser.open("http://127.0.0.1:{}/".format(port))
return proc
Usage
tb() # Starts a TensorBoard server on the logs directory, on port 6006
# and opens a new tab in your browser to use it.
tb("logs2", 6007) # Starts a second server on the logs2 directory, on port 6007,
# and opens a new tab to use it.
Starting a server does not block Jupyter (except for 2 seconds to ensure the server has the time to start before opening a tab). All TensorBoard servers will stop when you interrupt the kernel.
Advanced usage
If you want more control, you can kill the servers programmatically like this:
server1 = tb()
server2 = tb("logs2", 6007)
# and later...
server1.kill() # stops the first server
server2.kill() # stops the second server
You can set open_tab=False if you don't want new tabs to open. You can also set sleep to some other value if 2 seconds is too much or not enough on your system.
If you prefer to pause Jupyter while TensorBoard is running, then you can call any server's wait() method. This will block Jupyter until you interrupt the kernel, which will stop this server and all the others.
server1.wait()
Prerequisites
This solution assumes you have installed TensorBoard (e.g., using pip install tensorboard) and that it is available in the environment you started Jupyter in.
Acknowledgment
This answer was inspired by #SalvadorDali's answer. His solution is nice and simple, but I wanted to be able to start multiple tensorboard instances without blocking Jupyter. Also, I prefer not to delete log directories. Instead, I start tensorboard on the root log directory, and each TensorFlow run logs in a different subdirectory.

Answer 7

Another quick option with TF 2.x is through the plot_model() function. It's already built into more recent versions of TF utilities. For example:
import tensorflow
from tensorflow.keras.utils import plot_model
plot_model(model, to_file=('output_filename.png'))
This function is nice because you can have it display the layer name, output at a high DPI, configure it to plot horizontally, any other options. Here is the documentation for the function: https://www.tensorflow.org/api_docs/python/tf/keras/utils/plot_model
The plotting is very quick even for large models and works very well even with complex models that have multiple connections in and out.

Answer 8

TensorBoard Visualize Nodes - Architecture Graph
<img src="https://www.tensorflow.org/images/graph_vis_animation.gif" width=1300 height=680>