Simple way to visualize a TensorFlow graph in Jupy

问题:

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?

回答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('"', '&quot;'))
    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

回答2:

I wrote a Jupyter extension for tensorboard integration. It can:

1. Start tensorboard just by clicking a button in Jupyter
2. Manage multiple tensorboard instances.
3. Seamless integration with Jupyter interface.

Github: https://github.com/lspvic/jupyter_tensorboard

回答3:

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)

回答4:

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

回答5:

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.