TensorFlow
loss function
machine learning
neural networks
TensorFlow 2.0

Print all terms of loss function tensorflow 2.0

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

In TensorFlow 2, the total loss seen by the optimizer can include more than just the primary data loss. Regularization penalties, auxiliary losses, and custom terms may all contribute, so the clean way to “print all terms” is to compute them explicitly and log each piece before summing them.

Understand the Main Loss Versus Extra Losses

For a Keras model, the total loss often comes from two sources:

  • the main supervised loss, such as crossentropy or MSE
  • extra terms collected in model.losses, such as kernel regularizers

If you want visibility into every part, do not treat loss as a black box. Break it into named pieces yourself.

Inspect Regularization Terms with model.losses

Here is a small example:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential(
4    [
5        tf.keras.layers.Input(shape=(10,)),
6        tf.keras.layers.Dense(
7            16,
8            activation="relu",
9            kernel_regularizer=tf.keras.regularizers.l2(1e-4),
10        ),
11        tf.keras.layers.Dense(1),
12    ]
13)
14
15loss_fn = tf.keras.losses.MeanSquaredError()
16
17x = tf.random.normal((4, 10))
18y = tf.random.normal((4, 1))
19
20pred = model(x, training=True)
21main_loss = loss_fn(y, pred)
22reg_losses = model.losses
23total_loss = main_loss + tf.add_n(reg_losses) if reg_losses else main_loss
24
25print("main_loss:", float(main_loss.numpy()))
26for i, term in enumerate(reg_losses):
27    print(f"reg_loss_{i}:", float(term.numpy()))
28print("total_loss:", float(total_loss.numpy()))

This is the simplest way to expose the built-in terms that Keras is already tracking for you.

Use a Custom Training Step for Full Visibility

If you want this printed every step or every batch, a custom training loop or custom train_step is the right place:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential(
4    [
5        tf.keras.layers.Input(shape=(10,)),
6        tf.keras.layers.Dense(
7            16,
8            activation="relu",
9            kernel_regularizer=tf.keras.regularizers.l2(1e-4),
10        ),
11        tf.keras.layers.Dense(1),
12    ]
13)
14
15optimizer = tf.keras.optimizers.Adam()
16loss_fn = tf.keras.losses.MeanSquaredError()
17
18x = tf.random.normal((8, 10))
19y = tf.random.normal((8, 1))
20
21with tf.GradientTape() as tape:
22    pred = model(x, training=True)
23    main_loss = loss_fn(y, pred)
24    reg_losses = model.losses
25    reg_total = tf.add_n(reg_losses) if reg_losses else 0.0
26    total_loss = main_loss + reg_total
27
28print("main_loss:", float(main_loss.numpy()))
29print("reg_total:", float(reg_total.numpy() if hasattr(reg_total, "numpy") else reg_total))
30print("total_loss:", float(total_loss.numpy()))
31
32grads = tape.gradient(total_loss, model.trainable_variables)
33optimizer.apply_gradients(zip(grads, model.trainable_variables))

This makes the loss accounting explicit and easy to debug.

If your model includes several task-specific terms, store them in separate variables before summing:

python
1reconstruction_loss = tf.reduce_mean(tf.square(y - pred))
2consistency_loss = 0.1 * tf.reduce_mean(tf.abs(pred))
3reg_total = tf.add_n(model.losses) if model.losses else 0.0
4total_loss = reconstruction_loss + consistency_loss + reg_total
5
6print("reconstruction_loss:", float(reconstruction_loss.numpy()))
7print("consistency_loss:", float(consistency_loss.numpy()))
8print("reg_total:", float(reg_total.numpy() if hasattr(reg_total, "numpy") else reg_total))
9print("total_loss:", float(total_loss.numpy()))

That is usually better than trying to reverse-engineer a single scalar after the fact.

model.fit Can Log Terms Too

If you prefer model.fit, you can expose custom loss terms as metrics inside a subclassed model or custom callback. The important design idea is the same: calculate the terms separately, then report them separately.

Trying to “print all terms” without first structuring the loss computation into named variables is what makes the task awkward.

Common Pitfalls

The most common mistake is assuming model.losses already contains the main supervised loss. It does not. It usually contains only extra losses attached by layers or custom calls to add_loss.

Another pitfall is summing everything into one scalar before logging and then expecting to recover the components later. Once the values are merged, the individual contributions are gone unless you saved them first.

It is also easy to forget that some regularization terms appear only when the model has actually been called. If you inspect model.losses too early, it may be empty or incomplete.

Finally, printing tensors inside a traced function can behave differently from eager Python debugging. If you need reliable per-term inspection, start in eager mode or log through metrics and callbacks intentionally.

Summary

  • Break the loss into named components before adding them together.
  • The main supervised loss is separate from extra terms in model.losses.
  • Use a custom training step when you want full control over logging.
  • Regularization terms are easiest to inspect through model.losses.
  • If you want per-term visibility during training, structure the computation for observability instead of treating loss as one opaque scalar.

Related reading
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
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

All Rights Reserved.