如何在TensorFlow Estimator中每全局步获取训练损失并添加评估损失至TensorBoard？

要把评估损失也显示在TensorBoard的lossxx图表里，核心是让评估阶段也生成并保存对应的loss summary，并且和训练阶段的summary写入同一个目录，这样TensorBoard会自动合并同名的指标曲线。下面是具体的修改步骤和完整代码：

关键修改点 #

• 统一summary创建逻辑：把loss和accuracy的summary创建从训练分支抽出来，让训练和评估模式都能生成对应的summary
• 为评估模式添加SummarySaverHook：在评估模式的EstimatorSpec里指定evaluation_hooks，把评估阶段的summary保存到训练用的同一个目录
• 确保评估时执行summary保存：调用classifier.evaluate()来触发评估流程，写入评估summary

修改后的完整代码 #

class MyHook(tf.train.SessionRunHook):
    def after_run(self,run_context,run_value):
        _session = run_context.session
        _session.run(_session.graph.get_operation_by_name('acc_op'))

def my_model(features, labels, mode):
    # ... 保留你原有的网络结构代码 ...
    logits = tf.layers.dense(net, 3, activation=None)
    predicted_classes = tf.argmax(logits, 1)

    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'class': predicted_classes,
            'prob': tf.nn.softmax(logits)
        }
        return tf.estimator.EstimatorSpec(mode, predictions=predictions)

    # 计算损失和准确率，这部分和原代码一致
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    acc, acc_op = tf.metrics.accuracy(labels=labels, predictions=predicted_classes)
    tf.identity(acc_op,'acc_op')

    # 核心修改：把summary创建放在所有模式共享的区域，训练和评估都能生成
    loss_sum = tf.summary.scalar('lossxx', loss)
    accuracy_sum = tf.summary.scalar('accuracyxx', acc)
    merg = tf.summary.merge_all()

    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
        train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(
            mode, 
            loss=loss, 
            train_op=train_op,
            training_chief_hooks=[
                tf.train.SummarySaverHook(
                    save_steps=10, 
                    output_dir='./model', 
                    summary_op=merg
                )
            ]
        )

    # 核心修改：为评估模式添加SummarySaverHook，保存评估阶段的summary
    return tf.estimator.EstimatorSpec(
        mode, 
        loss=loss, 
        eval_metric_ops={'accuracy': (acc, acc_op)},
        evaluation_hooks=[
            tf.train.SummarySaverHook(
                save_steps=1,  # 评估通常一次跑完整验证集，所以每步保存一次即可
                output_dir='./model',  # 和训练目录一致，确保TensorBoard能合并数据
                summary_op=merg
            )
        ]
    )

# 训练代码不变
classifier.train(input_fn=train_input_fn, steps=1000,hooks=[ MyHook()])
# 新增：执行评估，触发评估summary的保存
classifier.evaluate(input_fn=eval_input_fn)

生效原理 #

• 训练阶段的training_chief_hooks会每隔10步把训练loss和accuracy写入./model目录
• 评估阶段的evaluation_hooks会在评估完成后把验证loss和accuracy写入同一个目录
• TensorBoard读取./model时，会自动把同名的lossxx（训练和评估的）合并到同一张图表里，x轴对齐全局步数

最后，重启TensorBoard（或者刷新页面）就能看到训练和评估的loss曲线同框展示了。

内容的提问来源于stack exchange，提问作者huaxz1986