TensorFlow v2: Graph does not appear on TensorBoard

[stephanwlee]

A graph may not appear if TensorFlow could not trace any graph during the execution. For instance, below code does not use @tf.function and, because it is pure eager, there is no graph TensorFlow used.

# Missing @tf.function
def foo(x, y):
  return tf.linalg.cross(x, y)

tf.summary.trace_on()
foo([1,2], [3,4])
with writer.as_default():
  tf.summary.trace_export()

Other known issues

• When a function is invoked before trace, the graph does not show up

@tf.function
def foo(x, y):
  return tf.linalg.cross(x, y)

foo([1,2], [3,4]) # this is not traced.
tf.summary.trace_on()
foo([1,2], [3,4])
with writer.as_default():
  tf.summary.trace_export()

[shashvatshahi1998]

So in both the cases graph does not appear??

[stephanwlee]

Latter is a bug of TensorFlow's and former is working as intended; if your code is purely eager, there is absolutely no computational graphs and thus the graph plugin cannot do anything about it.

[thealphacod3r]

@stephanwlee autograph can be used to bridge the gap between eager execution and graphs. It can be used to convert the eager-style python code to graph-generating code.

[cottrell]

So to clarify, if you use eager only code, you can not debug with tensorboard graphs? I'm confused as to how the keras writer does this. Maybe that is a tf.function under the hood now.

[stephanwlee]

@cottrell Yup correct. Pure eager builds no graph. Keras always has built some graph underneath (the details escape me at the moment but IIRC, in v1, it built a graph using its own session and, in v2, it constructs graph with FuncGraph) and is able to show some graph unless you pass run_eagerly to its compile method.

[WilledgeR]

I have the same problem

[simonschroder]

When using a model which inherits from tensorflow.python.keras.Model, there is a graph which can be visualized! This even works with run_eagerly=True. The trick is to save the graph at the right moment. Such a moment occurs, when keras/tensorflow builds the model, because it invokes model.call in non-eager mode (not exactly sure why). The resulting tensor holds the model graph (without losses afaik), which can be written using the "old" FileWriter. Example code that works for me with Tensorflow 1.13.1 and enabled eager execution:

import tensorflow
# Only use tensorflow's keras!
from tensorflow.python import keras as tfkeras
from tensorflow.python.training.rmsprop import RMSPropOptimizer
import numpy as np

tensorflow.enable_eager_execution()


class MyModel(tfkeras.Model):
    def __init__(self, tensorboard_folder_path):
        super(MyModel, self).__init__()
        self.dense1 = tfkeras.layers.LSTM(units=6)
        self.dense2 = tfkeras.layers.Dense(units=4)
        self.graph_has_been_written = False
        self.tensorboard_folder_path = tensorboard_folder_path

    def call(self, input, **kwargs):
        print("input shape", input.shape)
        result = self.dense1(input)
        result = self.dense2(result)
        if not tensorflow.executing_eagerly() and not self.graph_has_been_written:
            # In non eager mode and a graph is available which can be written to Tensorboard using the "old" FileWriter:
            model_graph = result.graph
            writer = tensorflow.summary.FileWriter(logdir=self.tensorboard_folder_path, graph=model_graph)
            writer.flush()
            self.graph_has_been_written = True
            print("Wrote eager graph to", self.tensorboard_folder_path)
        return result


if __name__ == "__main__":
    print("Eager execution:", tensorflow.executing_eagerly())
    # Create model and specify tensorboard folder:
    model = MyModel("/home/your_username/tensorboardtest/")
    optimizer = RMSPropOptimizer(learning_rate=0.001)
    model.compile(optimizer, tensorflow.losses.softmax_cross_entropy, run_eagerly=True)
    # Build the model (this will invoke model.call in non-eager mode). If model.build is not called explicitly here, it
    # will be called by model.fit_generator implicitly when the first batch is about to be feed to the network.
    model.build((None, None, 5))
    # Can only be called after the model has been built:
    model.summary()

    # Two arbitrary batches with different batch size and different sequence length:
    x1 = np.array([[[1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]],
                  dtype=np.float32)
    y1 = np.array([[1, 0, 0, 0]], dtype=np.float32)
    print("x1 shape", x1.shape)
    print("y1 shape", y1.shape)

    x2 = np.array([[[1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5]],
                   [[1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]], dtype=np.float32)
    y2 = np.array([[1, 0, 0, 0], [1, 0, 0, 0]], dtype=np.float32)
    print("x2 shape", x2.shape)
    print("y2 shape", y2.shape)

    # Simply yield the two batches alternately
    def iterator():
        switcher = False
        while 1:
            if switcher:
                yield x1, y1
            else:
                yield x2, y2
            switcher = not switcher

    model.fit_generator(iterator(), steps_per_epoch=10, epochs=1)

[tillns]

@simonschroder Unfortunately, this doesn't seem to work with tf 2.0 anymore. The use of FileWriter seems to have changed. Any ideas?

[luvwinnie]

Device: macbook pro 2018 , macOS Mojave 10.14.5
tensorflow version:
tensorflow==2.0.0a0
tensorflow-estimator-2.0-preview==1.14.0.dev2019052100

I have the same issues too, i followed the code shows in https://tensorflow.org/tensorboard/r2/graphs the most below method.

# The function to be traced.
@tf.function
def my_func(x, y):
  # A simple hand-rolled layer.
  return tf.nn.relu(tf.matmul(x, y))

# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = 'logs/func/%s' % stamp
writer = tf.summary.create_file_writer(logdir)

# Sample data for your function.
x = tf.random.uniform((3, 3))
y = tf.random.uniform((3, 3))

# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True, profiler=True)
# Call only one tf.function when tracing.
z = my_func(x, y)
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

in the tutorials, it shows the graphs, however when i ran this code in my environment it doesn't shows the graph in the tensorboard

[tillns]

I found out that one has to pass an input shape to the first layer of the model, so that keras is able to build the graph in advance. Otherwise, the tensorboard callback is not able to draw the graph because it lacks the predefined graph information from keras. I hope that I could help somebody with this information :)

[stephanwlee]

@luvwinnie I just tried your example and the graph was showing up correctly.

Especially because you put a timestamp onto the name of the logdir, I cannot imagine why you'd face a bug. Can you first try to reproduce this in very small minimal example? Thanks!

[stephanwlee]

I have made some incorrect statements about TF 2 and Keras so I would like to clear it up.

it built a graph using its own session and, in v2, it constructs graph with FuncGraph) and is able to show some graph unless you pass run_eagerly to its compile method.

Keras, in TF 2, even before calling compile or build, since layers are built with TensorFlow ops, can technically display a graph.

import tensorflow as tf
tensor = tf.keras.layers.Dense(32)(tf.keras.Input(32))
# https://github.com/tensorflow/tensorflow/blob/81bb13c6363f15581163976d298165e5bcff0588/tensorflow/python/framework/ops.py#L371-L374
tensor.graph # FuncGraph
tensor.graph.as_graph_def() # GraphDef

However, although I lack deep knowledge in this, it seemed like Keras makes few mutations to the graph and this is one of the reason why Keras implements [custom graph execution].(https://github.com/tensorflow/tensorflow/blob/81bb13c6363f15581163976d298165e5bcff0588/tensorflow/python/keras/backend.py#L3366).

You may be able to argue that some Keras subclass model can force Keras to spit out a graph but this statement is not generally true. Moreover, when user compiles Keras model to run eagerly, showing a graph will not help anyone as it can misrepresent how Keras will actually execute the code.

[olesalscheider]

Especially because you put a timestamp onto the name of the logdir, I cannot imagine why you'd face a bug. Can you first try to reproduce this in very small minimal example? Thanks!

I tried this example with TF2 from today's master branch and tensorboard 1.13.1. Here, tensorboard cannot find the exported graph...

[flobotics]

when following this https://www.tensorflow.org/tensorboard/r2/graphs with @tf.function it shows a "Graphs" -tab in tensorboard, but the loading bar only runs till the half and then stops loading, so there is no graph to see ?

I am using colab with

!pip3 install tensorflow-gpu==2.0.0-alpha0
!pip3 install tb-nightly