How do you decode one-hot labels in Tensorflow?

TensorFlow

one-hot encoding

machine learning

data preprocessing

neural networks

How do you decode one-hot labels 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 & Keras Machine Learning Data Science

Introduction

One-hot encoding converts a categorical label (like class 2 out of 5 classes) into a binary vector ([0, 0, 1, 0, 0]). Decoding reverses this: it takes the one-hot vector and returns the original class index. In TensorFlow, decoding is done with tf.argmax(), which returns the index of the maximum value along a specified axis. For model predictions (probability distributions), tf.argmax() picks the class with the highest predicted probability.

Basic Decoding with tf.argmax

python

1import tensorflow as tf
2
3# One-hot encoded labels (batch of 4 samples, 3 classes)
4one_hot = tf.constant([
5    [1, 0, 0],  # Class 0
6    [0, 1, 0],  # Class 1
7    [0, 0, 1],  # Class 2
8    [0, 1, 0],  # Class 1
9], dtype=tf.float32)
10
11# Decode: get the index of the 1 in each row
12labels = tf.argmax(one_hot, axis=1)
13print(labels.numpy())  # [0 1 2 1]

axis=1 means "find the max along columns for each row." Each row represents one sample, and the columns represent classes.

Decoding Model Predictions

Model outputs are probability distributions (from softmax), not clean one-hot vectors:

python

1# Simulated softmax output (probabilities sum to 1 per row)
2predictions = tf.constant([
3    [0.9, 0.05, 0.05],   # Confident: class 0
4    [0.1, 0.7, 0.2],     # Confident: class 1
5    [0.3, 0.3, 0.4],     # Less confident: class 2
6    [0.33, 0.34, 0.33],  # Very uncertain: class 1 (barely)
7])
8
9predicted_classes = tf.argmax(predictions, axis=1)
10print(predicted_classes.numpy())  # [0 1 2 1]
11
12# Get the confidence (max probability) for each prediction
13confidence = tf.reduce_max(predictions, axis=1)
14print(confidence.numpy())  # [0.9  0.7  0.4  0.34]

Mapping Indices to Class Names

python

1# Class name lookup
2class_names = ["cat", "dog", "bird"]
3
4# Using tf.gather
5indices = tf.argmax(one_hot, axis=1)
6names = tf.gather(class_names, indices)
7print(names.numpy())  # [b'cat' b'dog' b'bird' b'dog']
8
9# Using Python list indexing
10for i in indices.numpy():
11    print(class_names[i])
12# cat
13# dog
14# bird
15# dog

Decoding During Training and Evaluation

python

1import tensorflow as tf
2from tensorflow import keras
3
4# Build a simple model
5model = keras.Sequential([
6    keras.layers.Dense(64, activation='relu', input_shape=(784,)),
7    keras.layers.Dense(10, activation='softmax')
8])
9
10model.compile(
11    optimizer='adam',
12    loss='categorical_crossentropy',   # Use with one-hot labels
13    metrics=['accuracy']
14)
15
16# After training, decode predictions
17x_test = tf.random.normal((5, 784))
18predictions = model(x_test, training=False)  # Shape: (5, 10)
19
20# Decode to class indices
21predicted_classes = tf.argmax(predictions, axis=1)
22print(predicted_classes.numpy())  # e.g., [3 7 1 0 9]

One-Hot Encoding (The Reverse Operation)

python

1# Encode: integer labels to one-hot
2labels = tf.constant([0, 1, 2, 1])
3one_hot = tf.one_hot(labels, depth=3)
4print(one_hot.numpy())
5# [[1. 0. 0.]
6#  [0. 1. 0.]
7#  [0. 0. 1.]
8#  [0. 1. 0.]]
9
10# Decode: one-hot back to integer labels
11decoded = tf.argmax(one_hot, axis=1)
12print(decoded.numpy())  # [0 1 2 1]
13
14# Round-trip verification
15assert tf.reduce_all(tf.equal(labels, tf.cast(decoded, tf.int32)))

NumPy Equivalent

python

1import numpy as np
2
3one_hot = np.array([
4    [1, 0, 0],
5    [0, 1, 0],
6    [0, 0, 1],
7])
8
9labels = np.argmax(one_hot, axis=1)
10print(labels)  # [0 1 2]

The NumPy version works identically — np.argmax with axis=1.

Sparse vs Categorical Cross-Entropy

The choice between sparse_categorical_crossentropy and categorical_crossentropy determines whether you need one-hot encoding at all:

python

1# With one-hot labels: use categorical_crossentropy
2y_one_hot = tf.constant([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=tf.float32)
3model.compile(loss='categorical_crossentropy')
4
5# With integer labels: use sparse_categorical_crossentropy
6y_sparse = tf.constant([0, 1, 2])
7model.compile(loss='sparse_categorical_crossentropy')
8
9# sparse_categorical_crossentropy handles the one-hot conversion internally
10# This avoids the memory cost of storing large one-hot matrices

For large numbers of classes (e.g., 10,000+ in NLP), sparse labels save significant memory.

Multi-Label Decoding (Multiple Classes Per Sample)

If a sample can belong to multiple classes simultaneously, use thresholding instead of argmax:

python

1# Multi-label predictions (sigmoid output, not softmax)
2predictions = tf.constant([
3    [0.9, 0.1, 0.8],  # Classes 0 and 2
4    [0.2, 0.7, 0.6],  # Classes 1 and 2
5])
6
7# Threshold at 0.5
8multi_labels = tf.cast(predictions > 0.5, tf.int32)
9print(multi_labels.numpy())
10# [[1 0 1]
11#  [0 1 1]]

argmax only returns one index and is wrong for multi-label classification.

Batch Processing with tf.data

python

1# Decode labels in a tf.data pipeline
2def decode_one_hot(features, one_hot_labels):
3    return features, tf.argmax(one_hot_labels, axis=-1)
4
5dataset = tf.data.Dataset.from_tensor_slices((x_data, y_one_hot))
6dataset = dataset.map(decode_one_hot)
7dataset = dataset.batch(32)

Common Pitfalls

Wrong axis: tf.argmax(one_hot, axis=0) returns the max along rows (across samples), not within each sample. Use axis=1 for batch data where each row is a sample.
Using argmax for multi-label classification: argmax returns exactly one index. If samples can belong to multiple classes, use sigmoid activation with thresholding instead.
Ties in one-hot vectors: If two values are equal (e.g., [0.5, 0.5, 0]), argmax returns the first occurrence. This is nondeterministic for your purposes — consider whether ties matter in your application.
Forgetting to cast types: tf.argmax returns int64 by default. If your labels are int32, cast with tf.cast(decoded, tf.int32) before comparison.
Decoding logits vs probabilities: argmax works the same on raw logits (pre-softmax) and probabilities (post-softmax) because softmax preserves the relative ordering. You do not need to apply softmax before argmax.

Summary

Use tf.argmax(one_hot, axis=1) to decode one-hot labels to class indices
Works on both clean one-hot vectors and softmax probability distributions
Use tf.gather(class_names, indices) to map indices to human-readable names
Use sparse_categorical_crossentropy to avoid one-hot encoding entirely
For multi-label classification, use sigmoid + thresholding instead of argmax
argmax on logits gives the same result as argmax on probabilities — softmax is not needed for decoding

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

How do you decode one-hot labels in Tensorflow?

ML System Design practice on Codemia

Introduction

Basic Decoding with tf.argmax

Decoding Model Predictions

Mapping Indices to Class Names

Decoding During Training and Evaluation

One-Hot Encoding (The Reverse Operation)

NumPy Equivalent

Sparse vs Categorical Cross-Entropy

Multi-Label Decoding (Multiple Classes Per Sample)

Batch Processing with tf.data

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia