一个tf Estimator Summary Hook 函数

最新推荐文章于 2022-09-27 09:04:51 发布

原创最新推荐文章于 2022-09-27 09:04:51 发布 · 2k 阅读

0 ·

CC 4.0 BY-SA版权

tensorflow 专栏收录该内容

8 篇文章

订阅专栏

本文介绍了一种自定义的tfEstimatorSummaryHook函数，用于解决tf.estimator.DNNLinearCombinedClassifier中无法直接使用tensorboard查看loss下降曲线的问题。通过实现SessionRunHook，能够定期记录并展示训练过程中的损失变化。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一个tf Estimator Summary Hook 函数

应用场景：最近在使用tf.estimator.DNNLinearCombinedClassifier时，希望能够使用tensorboard去看一下loss的下降曲线。

虽然Tensorflow有提供了tf.train.SummarySaverHook，但是研究后发现，并不适用于tf.estimator.DNNLinearCombinedClassifier，只适用于自定义model_fn的时候。

于是经过研究，自己写了个hook函数，经过测试可以用。

class SummaryHook(tf.train.SessionRunHook):
        """ 
        Generate summary
        """
        def begin(self):
            self._step = -1
            self.summary_op = tf.summary.merge_all()
            
        def after_create_session(self, session, coord):
            self.summary_writer = tf.summary.FileWriter('./log_summary', session.graph)
            #self.summary_writer.add_summary(self.summary_op)

        def before_run(self, run_context):
            self._step += 1
            return tf.train.SessionRunArgs(self.summary_op)

        def after_run(self, run_context, run_values):
            flush_step = 50
            if self._step % flush_step == 0:
                print("Flusing summary")
                self.summary_writer.add_summary(run_values[0], global_step=self._step)
                self.summary_writer.flush()