learner_hooks.py 文件源码

def after_run(self, _run_context, run_values):
        if not self.is_chief or self._done:
            return

        step_done = run_values.results
        if self._active:
            log.info("Captured full trace at step %s", step_done)
            # Create output directory
            gfile.MakeDirs(self._output_dir)

            # Save run metadata
            trace_path = os.path.join(self._output_dir, "run_meta")
            with gfile.GFile(trace_path, "wb") as trace_file:
                trace_file.write(run_values.run_metadata.SerializeToString())
                log.info("Saved run_metadata to %s", trace_path)

            # Save timeline
            timeline_path = os.path.join(self._output_dir, "timeline.json")
            with gfile.GFile(timeline_path, "w") as timeline_file:
                tl_info = timeline.Timeline(run_values.run_metadata.step_stats)
                tl_chrome = tl_info.generate_chrome_trace_format(
                    show_memory=True)
                timeline_file.write(tl_chrome)
                log.info("Saved timeline to %s", timeline_path)

            # Save tfprof op log
            tf.profiler.write_op_log(
                graph=tf.get_default_graph(),
                log_dir=self._output_dir,
                run_meta=run_values.run_metadata)
            log.info("Saved op log to %s", self._output_dir)
            self._active = False
            self._done = True

        self._active = (step_done >= self.params["step"])