Practice Now
tensorflow
confusion matrix
machine learning
classification
tutorial

how to create confusion matrix for classification in tensorflow

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
TensorFlow & KerasMachine Learning

Introduction

A confusion matrix shows how a classification model's predictions compare to actual labels, breaking down true positives, false positives, true negatives, and false negatives for each class. In TensorFlow, use tf.math.confusion_matrix(labels, predictions) to compute it. For visualization, pass the result to seaborn.heatmap() or sklearn.metrics.ConfusionMatrixDisplay. The matrix helps identify which classes the model confuses most, guiding targeted improvements in data collection and model architecture.

Basic Confusion Matrix with TensorFlow

python
1import tensorflow as tf
2import numpy as np
3
4# True labels and model predictions
5y_true = [0, 1, 2, 0, 1, 2, 0, 1, 2]
6y_pred = [0, 1, 1, 0, 2, 2, 0, 1, 0]
7
8# Compute confusion matrix
9cm = tf.math.confusion_matrix(y_true, y_pred, num_classes=3)
10print(cm.numpy())
11# [[3 0 0]
12#  [0 2 1]
13#  [1 1 1]]

Each row represents the actual class and each column represents the predicted class. Row 0, column 0 (value 3) means all three class-0 samples were correctly predicted. Row 2, column 0 (value 1) means one class-2 sample was incorrectly predicted as class 0.

From Model Predictions to Confusion Matrix

python
1import tensorflow as tf
2import numpy as np
3
4# Assume a trained model
5model = tf.keras.Sequential([
6    tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
7    tf.keras.layers.Dense(3, activation='softmax')
8])
9model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
10
11# Generate predictions
12X_test = np.random.randn(100, 10)
13y_test = np.random.randint(0, 3, 100)
14
15# model.predict returns probabilities — convert to class labels
16y_pred_probs = model.predict(X_test)
17y_pred = np.argmax(y_pred_probs, axis=1)
18
19# Confusion matrix
20cm = tf.math.confusion_matrix(y_test, y_pred, num_classes=3)
21print(cm.numpy())

For multi-class classification with softmax output, apply np.argmax() to convert probability vectors to class indices before computing the confusion matrix.

Binary Classification

python
1import tensorflow as tf
2import numpy as np
3
4# Binary model with sigmoid output
5model = tf.keras.Sequential([
6    tf.keras.layers.Dense(32, activation='relu', input_shape=(5,)),
7    tf.keras.layers.Dense(1, activation='sigmoid')
8])
9
10X_test = np.random.randn(50, 5)
11y_test = np.array([0, 1, 1, 0, 1, 0, 0, 1] * 6 + [0, 1])
12
13# Sigmoid outputs probability — threshold at 0.5
14y_pred_probs = model.predict(X_test).flatten()
15y_pred = (y_pred_probs > 0.5).astype(int)
16
17cm = tf.math.confusion_matrix(y_test, y_pred, num_classes=2)
18print(cm.numpy())
19# [[TN, FP],
20#  [FN, TP]]

For binary classification, the 2x2 matrix positions are: top-left = True Negatives, top-right = False Positives, bottom-left = False Negatives, bottom-right = True Positives.

Visualizing with Seaborn

python
1import tensorflow as tf
2import numpy as np
3import seaborn as sns
4import matplotlib.pyplot as plt
5
6y_true = [0, 1, 2, 0, 1, 2, 0, 1, 2, 0]
7y_pred = [0, 1, 1, 0, 2, 2, 0, 1, 0, 0]
8class_names = ['Cat', 'Dog', 'Bird']
9
10cm = tf.math.confusion_matrix(y_true, y_pred).numpy()
11
12plt.figure(figsize=(8, 6))
13sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
14            xticklabels=class_names, yticklabels=class_names)
15plt.xlabel('Predicted')
16plt.ylabel('Actual')
17plt.title('Confusion Matrix')
18plt.tight_layout()
19plt.show()

sns.heatmap with annot=True displays the counts in each cell. Use fmt='d' for integer formatting or fmt='.2f' for normalized values.

Normalized Confusion Matrix

python
1import tensorflow as tf
2import numpy as np
3
4y_true = [0, 0, 0, 1, 1, 1, 1, 2, 2]
5y_pred = [0, 0, 1, 1, 1, 2, 1, 2, 0]
6
7cm = tf.math.confusion_matrix(y_true, y_pred).numpy()
8
9# Normalize by row (actual class) — shows recall per class
10cm_normalized = cm.astype('float') / cm.sum(axis=1, keepdims=True)
11print(np.round(cm_normalized, 2))
12# [[0.67 0.33 0.  ]
13#  [0.   0.75 0.25]
14#  [0.5  0.   0.5 ]]
15
16# Class 0: 67% correctly classified, 33% confused with class 1
17# Class 1: 75% correctly classified, 25% confused with class 2

Normalization reveals performance per class regardless of class imbalance. A raw count of 50 correct predictions for class A means different things if class A has 50 samples (100%) versus 500 samples (10%).

Computing Metrics from Confusion Matrix

python
1import tensorflow as tf
2import numpy as np
3
4y_true = [1, 0, 1, 1, 0, 1, 0, 0, 1, 0]
5y_pred = [1, 0, 1, 0, 0, 1, 1, 0, 1, 0]
6
7cm = tf.math.confusion_matrix(y_true, y_pred, num_classes=2).numpy()
8tn, fp, fn, tp = cm[0][0], cm[0][1], cm[1][0], cm[1][1]
9
10accuracy = (tp + tn) / (tp + tn + fp + fn)
11precision = tp / (tp + fp) if (tp + fp) > 0 else 0
12recall = tp / (tp + fn) if (tp + fn) > 0 else 0
13f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
14
15print(f"Accuracy:  {accuracy:.2f}")   # 0.80
16print(f"Precision: {precision:.2f}")   # 0.80
17print(f"Recall:    {recall:.2f}")      # 0.80
18print(f"F1 Score:  {f1:.2f}")          # 0.80

Using sklearn for Display (with TF Model)

python
1from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay, classification_report
2import matplotlib.pyplot as plt
3
4# After getting predictions from TF model
5y_true = [0, 1, 2, 0, 1, 2, 0, 1, 2]
6y_pred = [0, 1, 1, 0, 2, 2, 0, 1, 0]
7
8# sklearn provides formatted display
9cm = confusion_matrix(y_true, y_pred)
10disp = ConfusionMatrixDisplay(cm, display_labels=['Cat', 'Dog', 'Bird'])
11disp.plot(cmap='Blues', values_format='d')
12plt.title('Classification Results')
13plt.show()
14
15# Full classification report
16print(classification_report(y_true, y_pred, target_names=['Cat', 'Dog', 'Bird']))
17#               precision    recall  f1-score   support
18#          Cat       0.75      1.00      0.86         3
19#          Dog       0.67      0.67      0.67         3
20#         Bird       0.50      0.33      0.40         3

Common Pitfalls

  • Forgetting argmax for softmax output: model.predict() returns probability vectors for multi-class models. Passing raw probabilities to tf.math.confusion_matrix produces wrong results. Apply np.argmax(predictions, axis=1) first to convert to class indices.
  • Misinterpreting row vs column: In TensorFlow's confusion matrix, rows are actual labels and columns are predictions. Confusing this reverses the meaning of false positives and false negatives, leading to incorrect precision and recall calculations.
  • Ignoring class imbalance: A model predicting the majority class for all samples can show high accuracy but the confusion matrix reveals zero predictions for minority classes. Always examine per-class recall (diagonal divided by row sum) to catch this.
  • Using wrong num_classes: If num_classes is not specified, TensorFlow infers it from the data. If your test set happens to not contain class 4 out of 5 classes, the matrix will be 4x4 instead of 5x5. Always pass num_classes explicitly.
  • Not normalizing for comparison: Comparing raw confusion matrices across datasets with different sizes is misleading. A model tested on 1,000 samples naturally shows larger counts than one tested on 100. Normalize by row to get comparable recall rates.

Summary

  • Use tf.math.confusion_matrix(y_true, y_pred) to compute the confusion matrix in TensorFlow
  • Convert softmax probabilities to class indices with np.argmax() before computing the matrix
  • Rows represent actual labels, columns represent predicted labels
  • Visualize with seaborn.heatmap() or sklearn.metrics.ConfusionMatrixDisplay
  • Normalize by row to see per-class recall rates and identify confused class pairs
  • Extract precision, recall, and F1 directly from the matrix values (TP, FP, FN, TN)
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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.