Printing extra training metrics with Tensorflow Estimator

TensorFlow

Estimator

machine learning

training metrics

deep learning

Printing extra training metrics with Tensorflow Estimator

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Deep Learning TensorFlow & Keras Machine Learning

Introduction

TensorFlow's Estimator API logs loss by default during training, but you often need to monitor additional metrics like accuracy, learning rate, or custom values. The Estimator framework provides two mechanisms for this: tf.estimator.LoggingTensorHook to print tensor values during training, and tf.summary operations combined with TensorBoard for visualization. The key is adding the metrics to the EstimatorSpec and configuring hooks to display them.

Adding Metrics to the Model Function

The model function returns an EstimatorSpec that defines what to compute during training. Add extra metrics by creating tensors and passing them to a logging hook:

python

1import tensorflow as tf
2
3def model_fn(features, labels, mode):
4    # Build the model
5    logits = tf.keras.layers.Dense(10)(features['x'])
6    predictions = tf.argmax(logits, axis=1)
7
8    if mode == tf.estimator.ModeKeys.PREDICT:
9        return tf.estimator.EstimatorSpec(mode, predictions=predictions)
10
11    # Compute loss
12    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
13
14    # Compute extra metrics
15    accuracy = tf.metrics.accuracy(labels=labels, predictions=predictions)
16    precision = tf.metrics.precision(labels=labels, predictions=predictions)
17
18    if mode == tf.estimator.ModeKeys.EVAL:
19        return tf.estimator.EstimatorSpec(
20            mode, loss=loss,
21            eval_metric_ops={'accuracy': accuracy, 'precision': precision}
22        )
23
24    # Training mode — add logging hook
25    optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
26    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
27
28    # Create logging hook for training metrics
29    logging_hook = tf.estimator.LoggingTensorHook(
30        tensors={
31            'loss': loss,
32            'accuracy': accuracy[1],  # [1] is the update op
33            'step': tf.train.get_global_step()
34        },
35        every_n_iter=100  # Print every 100 steps
36    )
37
38    return tf.estimator.EstimatorSpec(
39        mode, loss=loss, train_op=train_op,
40        training_hooks=[logging_hook]
41    )

LoggingTensorHook Options

python

1# Print every N steps
2hook = tf.estimator.LoggingTensorHook(
3    tensors={'loss': loss, 'accuracy': accuracy_op},
4    every_n_iter=50
5)
6
7# Print every N seconds
8hook = tf.estimator.LoggingTensorHook(
9    tensors={'loss': loss, 'lr': learning_rate},
10    every_n_secs=30  # Every 30 seconds
11)
12
13# Custom formatting
14hook = tf.estimator.LoggingTensorHook(
15    tensors={'loss': loss, 'acc': accuracy_op},
16    every_n_iter=100,
17    formatter=lambda values: f"Step {values['global_step']}: loss={values['loss']:.4f}, acc={values['acc']:.3f}"
18)

Referencing Tensors by Name

If you cannot pass the tensor object directly, use the tensor's string name:

python

1# Name the tensor during graph construction
2accuracy_tensor = tf.identity(accuracy[1], name='training_accuracy')
3loss_tensor = tf.identity(loss, name='training_loss')
4
5# Reference by name in the hook
6hook = tf.estimator.LoggingTensorHook(
7    tensors={
8        'accuracy': 'training_accuracy:0',
9        'loss': 'training_loss:0'
10    },
11    every_n_iter=100
12)

This is useful when the hook is created outside the model function.

Adding TensorBoard Summaries

For persistent metric tracking, add summary operations:

python

1def model_fn(features, labels, mode):
2    # ... model definition ...
3
4    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
5
6    # Add summaries for TensorBoard
7    tf.summary.scalar('loss', loss)
8    tf.summary.scalar('accuracy', accuracy[1])
9    tf.summary.scalar('learning_rate', learning_rate)
10    tf.summary.histogram('logits', logits)
11
12    # Summaries are automatically written by the Estimator
13    # View them with: tensorboard --logdir=model_dir
14
15    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
16    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

bash

# Launch TensorBoard to view summaries
tensorboard --logdir=./model_dir --port=6006

Custom SessionRunHook

For more control, create a custom hook:

python

1class MetricPrinterHook(tf.estimator.SessionRunHook):
2    def __init__(self, tensors, every_n_steps=100):
3        self.tensors = tensors
4        self.every_n_steps = every_n_steps
5        self._step = 0
6
7    def before_run(self, run_context):
8        self._step += 1
9        if self._step % self.every_n_steps == 0:
10            return tf.estimator.SessionRunArgs(self.tensors)
11        return None
12
13    def after_run(self, run_context, run_values):
14        if run_values.results:
15            metrics = run_values.results
16            print(f"Step {self._step}: " +
17                  ", ".join(f"{k}={v:.4f}" for k, v in metrics.items()))
18
19# Use in model_fn
20hook = MetricPrinterHook(
21    tensors={'loss': loss, 'accuracy': accuracy_op},
22    every_n_steps=50
23)
24return tf.estimator.EstimatorSpec(
25    mode, loss=loss, train_op=train_op,
26    training_hooks=[hook]
27)

Setting Log Level

By default, TensorFlow only shows warnings. To see logging hook output:

python

1# Set log level to INFO
2tf.logging.set_verbosity(tf.logging.INFO)
3
4# Or via environment variable
5import os
6os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'  # Show all logs

Without setting the verbosity to INFO, LoggingTensorHook output is suppressed.

Complete Example

python

1import tensorflow as tf
2import numpy as np
3
4tf.logging.set_verbosity(tf.logging.INFO)
5
6def model_fn(features, labels, mode):
7    net = tf.keras.layers.Dense(128, activation='relu')(features['x'])
8    logits = tf.keras.layers.Dense(10)(net)
9    predictions = tf.argmax(logits, axis=1)
10
11    loss = tf.losses.sparse_softmax_cross_entropy(labels, logits)
12    accuracy = tf.metrics.accuracy(labels, predictions)
13
14    if mode == tf.estimator.ModeKeys.EVAL:
15        return tf.estimator.EstimatorSpec(
16            mode, loss=loss,
17            eval_metric_ops={'accuracy': accuracy}
18        )
19
20    optimizer = tf.train.AdamOptimizer(0.001)
21    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
22
23    hook = tf.estimator.LoggingTensorHook(
24        {'loss': loss, 'accuracy': accuracy[1]},
25        every_n_iter=100
26    )
27
28    return tf.estimator.EstimatorSpec(
29        mode, loss=loss, train_op=train_op,
30        training_hooks=[hook]
31    )
32
33# Train
34estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir='./model')
35estimator.train(input_fn=train_input_fn, steps=1000)
36# Output every 100 steps:
37# INFO:tensorflow:loss = 0.3245, accuracy = 0.8750

Common Pitfalls

Forgetting tf.logging.set_verbosity(tf.logging.INFO): Without this, LoggingTensorHook output is invisible. The hook runs but prints nothing because TF defaults to WARNING level.
Using accuracy[0] instead of accuracy[1]: tf.metrics.accuracy returns (value, update_op). Use index [1] (the update op) for training hooks because the value tensor is not updated during training without running the update op.
Passing hooks via train() instead of EstimatorSpec: Both work, but hooks in EstimatorSpec.training_hooks are model-specific, while hooks in estimator.train(hooks=[...]) are session-level. Use training_hooks for metrics tied to the model.
TF 2.x deprecation: The Estimator API is deprecated in TF 2.x. Use tf.keras with callbacks (tf.keras.callbacks.TensorBoard, custom Callback subclasses) for new projects.
Summary ops not appearing in TensorBoard: Summaries must be created inside the model function. Summaries created outside the Estimator's graph are not captured. Also ensure model_dir points to the correct directory.

Summary

Use tf.estimator.LoggingTensorHook to print extra metrics during training
Pass tensors as a dict to the hook and set every_n_iter or every_n_secs
Include the hook in EstimatorSpec(training_hooks=[hook])
Set tf.logging.set_verbosity(tf.logging.INFO) or output will be suppressed
Use tf.summary.scalar for TensorBoard visualization
For TF 2.x, prefer Keras callbacks instead of the Estimator API

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Printing extra training metrics with Tensorflow Estimator

ML System Design practice on Codemia

Introduction

Adding Metrics to the Model Function

LoggingTensorHook Options

Referencing Tensors by Name

Adding TensorBoard Summaries

Custom SessionRunHook

Setting Log Level

Complete Example

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia