How can I use tensorboard with tf.estimator.Estimator

Question

I am considering to move my code base to tf.estimator.Estimator, but I cannot find an example on how to use it in combination with tensorboard summaries.

MWE:

<

Answer 1

EDIT: Upon testing (in v1.1.0, and probably in later versions as well), it is apparent that tf.estimator.Estimator will automatically write summaries for you. I confirmed this using OP's code and tensorboard.

(Some poking around r1.4 leads me to conclude that this automatic summary writing occurs due to tf.train.MonitoredTrainingSession.)

Ultimately, the automatic summarizing is accomplished with the use of hooks, so if you wanted to customize the Estimator's default summarizing, you could do so using hooks. Below are the (edited) details from the original answer.

---

You'll want to use hooks, formerly known as monitors. (Linked is a conceptual/quickstart guide; the short of it is that the notion of hooking into / monitoring training is built into the Estimator API. A bit confusingly, though, it doesn't seem like the deprecation of monitors for hooks is really documented except in a deprecation annotation in the actual source code...)

Based on your usage, it looks like r1.2's SummarySaverHook fits your bill.

summary_hook = tf.train.SummarySaverHook(
    SAVE_EVERY_N_STEPS,
    output_dir='/tmp/tf',
    summary_op=tf.summary.merge_all())

You may want to customize the hook's initialization parameters, as by providing an explicity SummaryWriter or writing every N seconds instead of N steps.

If you pass this into the EstimatorSpec, you'll get your customized Summary behavior:

return tf.estimator.EstimatorSpec(mode=mode, predictions=y,loss=loss,
                                  train_op=train,
                                  training_hooks=[summary_hook])

EDIT NOTE: A previous version of this answer suggested passing the summary_hook into estimator.train(input_fn=input_fn, steps=5, hooks=[summary_hook]). This does not work because tf.summary.merge_all() has to be called in the same context as your model graph.