如何在分布式环境中使用 Estimator API 在 Tensorboard 中显示运行时统计信息答案

【问题标题】：How to display Runtime Statistics in Tensorboard using Estimator API in a distributed environment如何在分布式环境中使用 Estimator API 在 Tensorboard 中显示运行时统计信息
【发布时间】：2018-01-24 22:16:45
【问题描述】：

This article 说明了如何将运行时统计信息添加到 Tensorboard：

    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    summary, _ = sess.run([merged, train_step],
                          feed_dict=feed_dict(True),
                          options=run_options,
                          run_metadata=run_metadata)
    train_writer.add_run_metadata(run_metadata, 'step%d' % i)
    train_writer.add_summary(summary, i)
    print('Adding run metadata for', i)

在 Tensorboard 中创建以下细节：

这在一台机器上相当简单。如何在分布式环境中使用 Estimator 做到这一点？

【问题讨论】：

我也遇到了同样的问题，你找到解决办法了吗？
钩子可能是这样的
还没有。有关如何使用钩子执行此操作的任何指示？

标签： python tensorflow tensorboard

【解决方案1】：

您可以使用tf.train.ProfilerHook。然而问题是它是在 1.14 发布的。

示例用法：

estimator = tf.estimator.LinearClassifier(...)
hooks = [tf.train.ProfilerHook(output_dir=model_dir, save_secs=600, show_memory=False)]
estimator.train(input_fn=train_input_fn, hooks=hooks)

执行钩子会在output_dir中生成文件timeline-xx.json。

然后在 chrome 浏览器中打开 chrome://tracing/ 并加载文件。您将获得如下所示的时间使用时间表。

【讨论】：

【解决方案2】：

我使用以下基于 ProfilerHook 的钩子让估算器将运行元数据输出到模型目录中，然后使用 Tensorboard 对其进行检查。

import tensorflow as tf
from tensorflow.python.training.session_run_hook import SessionRunHook, SessionRunArgs
from tensorflow.python.training import training_util
from tensorflow.python.training.basic_session_run_hooks import SecondOrStepTimer

class MetadataHook(SessionRunHook):
    def __init__ (self,
                  save_steps=None,
                  save_secs=None,
                  output_dir=""):
        self._output_tag = "step-{}"
        self._output_dir = output_dir
        self._timer = SecondOrStepTimer(
            every_secs=save_secs, every_steps=save_steps)

    def begin(self):
        self._next_step = None
        self._global_step_tensor = training_util.get_global_step()
        self._writer = tf.summary.FileWriter (self._output_dir, tf.get_default_graph())

        if self._global_step_tensor is None:
            raise RuntimeError("Global step should be created to use ProfilerHook.")

    def before_run(self, run_context):
        self._request_summary = (
            self._next_step is None or
            self._timer.should_trigger_for_step(self._next_step)
        )
        requests = {"global_step": self._global_step_tensor}
        opts = (tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
            if self._request_summary else None)
        return SessionRunArgs(requests, options=opts)

    def after_run(self, run_context, run_values):
        stale_global_step = run_values.results["global_step"]
        global_step = stale_global_step + 1
        if self._request_summary:
            global_step = run_context.session.run(self._global_step_tensor)
            self._writer.add_run_metadata(
                run_values.run_metadata, self._output_tag.format(global_step))
            self._writer.flush()
        self._next_step = global_step + 1

    def end(self, session):
        self._writer.close()

要使用它，需要像往常一样创建估算器实例 (my_estimator)，无论它是预制估算器还是自定义估算器。所需的操作称为将上述类的实例作为挂钩传递。例如：

hook = MetadataHook(save_steps=1, output_dir=<model dir>)
my_estimator.train( train_input_fn, hooks=[hook] )

运行元数据将放置在模型目录中，并且可以通过 TensorBoard 进行检查。

【讨论】：