Using the Tensorflow profiler with tf.Estimator

Question

I need to use the Tensorflow profiler to profile some code that is running slowly for some reason. Unfortunately, the code in question uses tf.Estimator, so I can't figure out how to inject the run metadata object into the session run() call in order to get the information that the profiler needs.

What should I do?

Latex

Answer 1

tf.estimator use tf.train.ProfilerHook works!

just add a ProfilerHook in TrainSpec hooks!

hook = tf.train.ProfilerHook(
    save_steps=20,
    output_dir=os.path.join(args.model_dir, "tracing"),
    show_dataflow=True,
    show_memory=True)
hooks = [hook]
train_spec = tf.estimator.TrainSpec(
    hooks=hooks,
    input_fn=lambda: input_fn())

then, you could get the tracing file like timeline-{}.json in model_dir/tracing, and open chrome chrome://tracing to visual!

refs: https://mcmap.net/q/916316/-how-to-display-runtime-statistics-in-tensorboard-using-estimator-api-in-a-distributed-environment

Counterstatement

Answer 2

with tf.contrib.tfprof.ProfileContext('/tmp/train_dir', dump_steps=[10]) as pctx: estimator.train() # any thing you want to profile

Then you will get a file at /tmp/train_dir/profile_10

Arguments are defined in https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/profiler/profile_context.py

Dynamite

Answer 3

Use ProfileContext, as described here: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler . This allows you to profile without needing to get access to the session.

Latex