TensorFlow
Confusion Matrix
TensorBoard
Machine Learning
Data Visualization

Tensorflow Confusion Matrix in TensorBoard

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

A confusion matrix is one of the most useful ways to inspect classification errors because it shows exactly which classes are being confused with each other. TensorBoard does not create it automatically for every model, but TensorFlow makes it straightforward to compute the matrix and log it as an image or summary artifact during evaluation.

Compute the Matrix From Labels and Predictions

The first step is to collect the true labels and predicted labels for a validation set. In TensorFlow, the raw confusion matrix itself is easy to compute with tf.math.confusion_matrix.

python
1import tensorflow as tf
2
3true_labels = tf.constant([0, 1, 2, 1, 0, 2])
4pred_labels = tf.constant([0, 2, 2, 1, 0, 1])
5
6cm = tf.math.confusion_matrix(true_labels, pred_labels, num_classes=3)
7print(cm.numpy())

The rows represent true classes and the columns represent predicted classes. Diagonal values are correct predictions. Off-diagonal values are mistakes.

That raw tensor is already useful for debugging, but TensorBoard becomes much more informative if you visualize it clearly.

Log the Matrix to TensorBoard as an Image

A common pattern is to convert the confusion matrix into a matplotlib figure and write it with tf.summary.image.

python
1import io
2import matplotlib.pyplot as plt
3import numpy as np
4import tensorflow as tf
5
6
7def plot_confusion_matrix(cm, class_names):
8    fig, ax = plt.subplots(figsize=(4, 4))
9    ax.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
10    ax.set_xticks(np.arange(len(class_names)))
11    ax.set_yticks(np.arange(len(class_names)))
12    ax.set_xticklabels(class_names, rotation=45, ha="right")
13    ax.set_yticklabels(class_names)
14    ax.set_xlabel("Predicted")
15    ax.set_ylabel("True")
16
17    for i in range(cm.shape[0]):
18        for j in range(cm.shape[1]):
19            ax.text(j, i, cm[i, j], ha="center", va="center", color="black")
20
21    fig.tight_layout()
22    return fig
23
24
25def figure_to_image(fig):
26    buf = io.BytesIO()
27    fig.savefig(buf, format="png")
28    plt.close(fig)
29    buf.seek(0)
30    image = tf.image.decode_png(buf.getvalue(), channels=4)
31    return tf.expand_dims(image, 0)
32
33
34cm = np.array([[2, 0, 0], [0, 1, 1], [0, 1, 1]])
35file_writer = tf.summary.create_file_writer("logs/cm")
36
37with file_writer.as_default():
38    fig = plot_confusion_matrix(cm, ["cat", "dog", "bird"])
39    tf.summary.image("confusion_matrix", figure_to_image(fig), step=1)

When TensorBoard reads that log directory, the matrix appears in the Images tab.

Integrate It With Model Evaluation

In practice, you usually compute the matrix after running inference on validation data. A simple Keras workflow looks like this:

python
1import numpy as np
2
3probs = model.predict(x_val, verbose=0)
4preds = np.argmax(probs, axis=1)
5true = np.argmax(y_val, axis=1)
6cm = tf.math.confusion_matrix(true, preds, num_classes=num_classes).numpy()

That output can be logged at the end of each epoch or after periodic evaluation. Logging every single training step is rarely worth it because confusion matrices are evaluation summaries, not high-frequency training metrics.

Normalization is also often helpful. If one class is much larger than the others, raw counts can hide the pattern. Dividing each row by its sum shows error proportions more clearly.

Why TensorBoard Helps Here

Scalar accuracy can look healthy while one class is failing badly. A confusion matrix reveals those asymmetries immediately.

For example, if classes 1 and 2 are repeatedly confused, the matrix shows strong off-diagonal counts even when overall accuracy seems acceptable. That often points to:

  • class imbalance
  • poor labels
  • missing features
  • too-similar classes
  • threshold or calibration issues

That is why the confusion matrix is best treated as a diagnostic view, not just a pretty chart.

Common Pitfalls

A common mistake is mixing integer class labels with one-hot labels without converting them consistently before computing the matrix.

Another mistake is computing the matrix on training data and using it as if it represented generalization. The validation or test set is the more informative source.

People also often forget to label the axes with class names. A matrix with unlabeled rows and columns is much less useful when the model has more than two classes.

Finally, be careful with normalization. Normalized matrices are great for pattern detection, but raw counts still matter for understanding class frequency.

Summary

  • Compute the raw matrix with tf.math.confusion_matrix.
  • Convert it into an image if you want to inspect it inside TensorBoard.
  • Log it after evaluation, not necessarily at every training step.
  • Use validation predictions so the matrix reflects real model behavior.
  • Class names and optional normalization make the visualization much easier to interpret.

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.