Migrating tf.summary usage to TF 2.0

View on TensorFlow.org

Run in Google Colab

View source on GitHub

!pip install -q tf-nightly-2.0-preview

TensorFlow 2.x selected.

TensorFlow 2.0 includes significant changes to the tf.summary API used to write summary data for visualization in TensorBoard.

What's changed

It's useful to think of the tf.summary API as two sub-APIs:

• A set of ops for recording individual summaries - summary.scalar(), summary.histogram(), summary.image(), summary.audio(), and summary.text() - which are called inline from your model code.
• Writing logic that collects these individual summaries and writes them to a specially formatted log file (which TensorBoard then reads to generate visualizations).

In TF 1.x

The two halves had to be manually wired together - by fetching the summary op outputs via Session.run() and calling FileWriter.add_summary(output, step). The v1.summary.merge_all() op made this easier by using a graph collection to aggregate all summary op outputs, but this approach still worked poorly for eager execution and control flow, making it especially ill-suited for TF 2.0.

In TF 2.X

The two halves are tightly integrated, and now individual tf.summary ops write their data immediately when executed. Using the API from your model code should still look familiar, but it's now friendly to eager execution while remaining graph-mode compatible. Integrating both halves of the API means the summary.FileWriter is now part of the TensorFlow execution context and gets accessed directly by tf.summary ops, so configuring writers is the main part that looks different.

Example usage with eager execution, the default in TF 2.0:

writer = tf.summary.create_file_writer("/tmp/mylogs/eager")

with writer.as_default():
  for step in range(100):
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)
    writer.flush()

ls /tmp/mylogs/eager

Example usage with tf.function graph execution:

writer = tf.summary.create_file_writer("/tmp/mylogs/tf_function")

@tf.function
def my_func(step):
  with writer.as_default():
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)

for step in range(100):
  my_func(step)
  writer.flush()

ls /tmp/mylogs/tf_function

Example usage with legacy TF 1.x graph execution:

g = tf.compat.v1.Graph()
with g.as_default():
  step = tf.Variable(0, dtype=tf.int64)
  step_update = step.assign_add(1)
  writer = tf.summary.create_file_writer("/tmp/mylogs/session")
  with writer.as_default():
    tf.summary.scalar("my_metric", 0.5, step=step)
  all_summary_ops = tf.compat.v1.summary.all_v2_summary_ops()
  writer_flush = writer.flush()


with tf.compat.v1.Session(graph=g) as sess:
  sess.run([writer.init(), step.initializer])

  for i in range(100):
    sess.run(all_summary_ops)
    sess.run(step_update)
    sess.run(writer_flush)  

ls /tmp/mylogs/session

Converting your code

Converting existing tf.summary usage to the TF 2.0 API cannot be reliably automated, so the tf_upgrade_v2 script just rewrites it all to tf.compat.v1.summary. To migrate to TF 2.0, you'll need to adapt your code as follows:

1. A default writer set via .as_default() must be present to use summary ops

   • This means executing ops eagerly or using ops in graph construction
   • Without a default writer, summary ops become silent no-ops
   • Default writers do not (yet) propagate across the @tf.function execution boundary - they are only detected when the function is traced - so best practice is to call writer.as_default() within the function body, and to ensure that the writer object continues to exist as long as the @tf.function is being used

2. The "step" value must be passed into each op via a the step argument

   • TensorBoard requires a step value to render the data as a time series
   • Explicit passing is necessary because the global step from TF 1.x has been removed, so each op must know the desired step variable to read
   • To reduce boilerplate, experimental support for registering a default step value is available as tf.summary.experimental.set_step(), but this is provisional functionality that may be changed without notice

3. Function signatures of individual summary ops have changed

   • Return value is now a boolean (indicating if a summary was actually written)
   • The second parameter name (if used) has changed from tensor to data
   • The collections parameter has been removed; collections are TF 1.x only
   • The family parameter has been removed; just use tf.name_scope()

4. [Only for legacy graph mode / session execution users]

   • First initialize the writer with v1.Session.run(writer.init())

   • Use v1.summary.all_v2_summary_ops() to get all TF 2.0 summary ops for the current graph, e.g. to execute them via Session.run()

   • Flush the writer with v1.Session.run(writer.flush()) and likewise for close()

If your TF 1.x code was instead using tf.contrib.summary API, it's much more similar to the TF 2.0 API, so tf_upgrade_v2 script will automate most of the migration steps (and emit warnings or errors for any usage that cannot be fully migrated). For the most part it just rewrites the API calls to tf.compat.v2.summary; if you only need compatibility with TF 2.0+ you can drop the compat.v2 and just reference it as tf.summary.

Additional tips

In addition to the critical areas above, some auxiliary aspects have also changed:

• Conditional recording (like "log every 100 steps") has a new look

  • To control ops and associated code, wrap them in a regular if statement (which works in eager mode and in @tf.function via autograph) or a tf.cond
  • To control just summaries, use the new tf.summary.record_if() context manager, and pass it the boolean condition of your choosing
  • These replace the TF 1.x pattern:

if condition:
  writer.add_summary()

• No direct writing of tf.compat.v1.Graph - instead use trace functions

  • Graph execution in TF 2.0 uses @tf.function instead of the explicit Graph
  • In TF 2.0, use the new tracing-style APIs tf.summary.trace_on() and tf.summary.trace_export() to record executed function graphs

• No more global writer caching per logdir with tf.summary.FileWriterCache

  • Users should either implement their own caching/sharing of writer objects, or just use separate writers (TensorBoard support for the latter is in progress)

• The event file binary representation has changed

  • TensorBoard 1.x already supports the new format; this difference only affects users who are manually parsing summary data from event files
  • Summary data is now stored as tensor bytes; you can use tf.make_ndarray(event.summary.value[0].tensor) to convert it to numpy