TensorFlow
custom metrics
machine learning
NaN error
streaming metrics

Custom metric based on tensorflow's streaming metrics returns NaN

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

Introduction

When a custom TensorFlow metric built from streaming metrics returns NaN, the problem is usually not the formula itself. The usual causes are reading the metric value before its update op has executed, dividing by zero during early batches, or forgetting to initialize the local variables that store streaming state.

Why streaming metrics behave differently

The old TensorFlow streaming metric API, now mostly seen through tf.compat.v1.metrics.*, returns two things:

  1. A value tensor that reads the current accumulated state.
  2. An update op that changes that state using the latest batch.

That means the value tensor is not automatically refreshed just because you built it. If your custom metric combines value tensors without controlling update order, you can easily compute a ratio from stale zeros, which leads to NaN.

For example, an F1 score derived from precision and recall is vulnerable because the denominator is precision + recall. Early in evaluation, both may still be zero.

A correct TensorFlow 1 style pattern

The safe approach is to run the update ops first and to use divide_no_nan for the final ratio:

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5labels = tf.constant([1, 0, 1, 1], dtype=tf.int32)
6predictions = tf.constant([1, 0, 0, 1], dtype=tf.int32)
7
8precision, precision_update = tf.compat.v1.metrics.precision(labels, predictions)
9recall, recall_update = tf.compat.v1.metrics.recall(labels, predictions)
10
11with tf.control_dependencies([precision_update, recall_update]):
12    f1 = tf.math.divide_no_nan(2.0 * precision * recall, precision + recall)
13
14with tf.compat.v1.Session() as sess:
15    sess.run(tf.compat.v1.global_variables_initializer())
16    sess.run(tf.compat.v1.local_variables_initializer())
17    print(sess.run(f1))

There are three important details in that example:

  1. precision_update and recall_update run before f1 is read.
  2. tf.math.divide_no_nan prevents a 0 / 0 result from becoming NaN.
  3. local_variables_initializer() is required because streaming metrics store internal counts in local variables.

If you skip any of those steps, the metric may remain invalid even though the graph builds successfully.

What goes wrong in custom formulas

Suppose you build a custom metric like this:

python
precision, precision_update = tf.compat.v1.metrics.precision(labels, predictions)
recall, recall_update = tf.compat.v1.metrics.recall(labels, predictions)
f1 = 2.0 * precision * recall / (precision + recall)

This looks reasonable, but it has two hidden problems. First, the formula reads the value tensors directly without forcing the update ops to execute. Second, ordinary division will return NaN when both precision and recall are zero, which is common at startup or on edge-case batches.

A lot of confusion comes from the fact that the metric variables are stateful while the formula looks like ordinary stateless tensor math. Streaming metrics are accumulators, so execution order matters.

Modern TensorFlow 2 alternative

If you are in TensorFlow 2 or Keras, the cleaner approach is often to write a tf.keras.metrics.Metric subclass and keep the state yourself:

python
1import tensorflow as tf
2
3class BinaryF1(tf.keras.metrics.Metric):
4    def __init__(self, name="binary_f1", **kwargs):
5        super().__init__(name=name, **kwargs)
6        self.tp = self.add_weight(name="tp", initializer="zeros")
7        self.fp = self.add_weight(name="fp", initializer="zeros")
8        self.fn = self.add_weight(name="fn", initializer="zeros")
9
10    def update_state(self, y_true, y_pred, sample_weight=None):
11        y_true = tf.cast(y_true, tf.float32)
12        y_pred = tf.cast(y_pred > 0.5, tf.float32)
13        self.tp.assign_add(tf.reduce_sum(y_true * y_pred))
14        self.fp.assign_add(tf.reduce_sum((1.0 - y_true) * y_pred))
15        self.fn.assign_add(tf.reduce_sum(y_true * (1.0 - y_pred)))
16
17    def result(self):
18        precision = tf.math.divide_no_nan(self.tp, self.tp + self.fp)
19        recall = tf.math.divide_no_nan(self.tp, self.tp + self.fn)
20        return tf.math.divide_no_nan(2.0 * precision * recall, precision + recall)
21
22    def reset_state(self):
23        self.tp.assign(0.0)
24        self.fp.assign(0.0)
25        self.fn.assign(0.0)

This avoids the awkward separation between value tensors and update ops and is usually easier to reason about in modern projects.

Common Pitfalls

The most common mistake is forgetting tf.compat.v1.local_variables_initializer(). Streaming metrics depend on local variables, so initializing only global variables is not enough.

Another issue is reading precision or recall before running their update ops. If the counts are still zero, the formula can produce invalid results or misleading zeros.

Division by zero is also common on highly imbalanced batches or empty positive classes. Use tf.math.divide_no_nan anywhere the denominator might legitimately be zero.

Finally, avoid mixing TensorFlow 1 streaming metrics into TensorFlow 2 Keras code unless you have a compatibility reason. A native Metric subclass is usually clearer and less error-prone.

Summary

  • Streaming metrics return both a value tensor and an update op, and the update order matters.
  • 'NaN usually comes from stale zero state, missing local variable initialization, or ordinary division by zero.'
  • In TensorFlow 1 style code, use control dependencies and tf.math.divide_no_nan.
  • In TensorFlow 2, prefer a custom tf.keras.metrics.Metric subclass.
  • Treat metric code as stateful logic, not just a pure formula.

Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

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

Start the free course
Basis Lab
No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my course
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