Debugging, profiling and export formats

Shape and dtype errors, NaN losses, the tf.debugging toolkit, the profiler, and exporting to TF Lite, TF.js and TF Serving from one model.

The errors you will actually hit

import tensorflow as tf

# asserts that run in the graph and raise with a readable message
tf.debugging.assert_shapes([
    (x, ("batch", "features")),
    (y, ("batch",)),
])

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        tf.debugging.assert_all_finite(logits, "non-finite logits")
        loss = loss_fn(y, logits)
        tf.debugging.assert_all_finite(loss, "non-finite loss")
    grads = tape.gradient(loss, model.trainable_variables)
    # check_numerics surfaces the offending op instead of letting NaN spread
    grads = [tf.debugging.check_numerics(g, f"grad {v.name}")
             for g, v in zip(grads, model.trainable_variables)]
    optimiser.apply_gradients(zip(grads, model.trainable_variables))
    return loss

• Debug in eager mode first (model(x) outside a tf.function). Eager errors point at the exact line, graph errors point at a node name.
• Use tf.debugging.enable_check_numerics() to make the first non-finite value raise at its source rather than silently propagate to the loss.
• Run the model on a single tiny batch and try to overfit it to zero loss. If that fails, the bug is structural; if it succeeds, the bug is in the data or the schedule.

Profiling a training step

# the profiler records a short window; open the result in TensorBoard
tf.profiler.experimental.start("logs/profile")

for step, (x, y) in enumerate(train_ds.take(20)):
    with tf.profiler.experimental.Trace("train", step_num=step):
        loss = train_step(x, y)
    if step == 10:
        break

tf.profiler.experimental.stop()   # then: tensorboard --logdir logs/profile

• Read the Trace Viewer top-down: the largest bar is the bottleneck. If the input pipeline bar is empty much of the time, you are input-bound.
• Profile short windows and warm up first. A profile of the first ten steps measures graph construction, not steady-state training.
• Compare kernel time against total step time. A large gap means the accelerator is waiting, and the fix is almost always in the input pipeline.
• Turn off profiling for the real run: it adds overhead and produces large files.

Export formats and their trade-offs

import tensorflow as tf

# 1. Keras format: the full model with preprocessing layers, for resuming work
model.save("model.keras")

# 2. SavedModel: the serving format, with a fixed inference signature
class ServingModule(tf.Module):
    def __init__(self, model):
        self.model = model

    @tf.function(input_signature=[
        tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32, name="image")])
    def serve(self, image):
        return {"scores": tf.nn.softmax(self.model(image, training=False))}

module = ServingModule(model)
tf.saved_model.save(module, "saved_model",
                    signatures={"serving_default": module.serve})

# 3. TF Lite: quantised, for mobile and edge
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]    # int8 post-training quantisation
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
tflite_model = converter.convert()
open("model.tflite", "wb").write(tflite_model)

FAQ

Related

TensorBoard and experiment tracking Mixed precision and multi-device training

Last refreshed 2026-09-18.