tf.profiler.experimental.Trace

Context manager that generates a trace event in the profiler.

tf.profiler.experimental.Trace(
    name, **kwargs
)

A trace event will start when entering the context, and stop and save the result to the profiler when exiting the context. Open TensorBoard Profile tab and choose trace viewer to view the trace event in the timeline.

Trace events are created only when the profiler is enabled. More information on how to use the profiler can be found at https://tensorflow.org/guide/profiler

Example usage:

tf.profiler.experimental.start('logdir')
for step in range(num_steps):
  # Creates a trace event for each training step with the step number.
  with tf.profiler.experimental.Trace("Train", step_num=step):
    train_fn()
tf.profiler.experimental.stop()

tf.profiler.experimental.start('logdir')
for step in range(num_steps):
# Creates a trace event for each training step with the
# step number.
with tf.profiler.experimental.Trace("Train", step_num=step):
train_fn()
tf.profiler.experimental.stop()

Methods

set_metadata

View source

set_metadata(
    **kwargs
)

Sets metadata in this trace event.

This method enables setting metadata in a trace event after it is created.

Example usage:

def call(function):
  with tf.profiler.experimental.Trace("call",
       function_name=function.name) as tm:
    binary, in_cache = jit_compile(function)
    tm.set_metadata(in_cache=in_cache)
    execute(binary)

In this example, we want to trace how much time spent on calling a function, which includes compilation and execution. The compilation can be either getting a cached copy of the binary or actually generating the binary, which is indicated by the boolean "in_cache" returned by jit_compile(). We need to use set_metadata() to pass in_cache because we did not know the in_cache value when the trace was created (and we cannot create the trace after jit_compile(), because we want to measure the entire duration of call()).

__enter__

View source

__enter__()

__exit__

View source

__exit__(
    exc_type, exc_val, exc_tb
)