Softmax output layer in TensorFlow

Introduction

The softmax output layer helps a model pick one choice from many by turning numbers into probabilities that add up to 1.

When you want the model to classify images into multiple categories like cats, dogs, or birds.

When predicting the next word in a sentence from a list of possible words.

When sorting emails into folders like work, personal, or spam.

When deciding which product a customer might buy from many options.

When recognizing handwritten digits from 0 to 9.

Syntax

TensorFlow

model.add(tf.keras.layers.Dense(number_of_classes, activation='softmax'))

The number_of_classes is how many categories you want to predict.

Softmax turns the output into probabilities that sum to 1.

Examples

This creates a softmax layer for 3 classes, like red, green, or blue.

TensorFlow

model.add(tf.keras.layers.Dense(3, activation='softmax'))

A simple model with one softmax output layer for 10 classes.

TensorFlow

model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='softmax')
])

Sample Model

This code builds a small neural network with a softmax output layer for 3 classes. It trains on some simple data and then predicts the probabilities for one new example. The output shows the probabilities for each class and confirms they add up to 1.

TensorFlow

import tensorflow as tf
import numpy as np

# Create a simple model with softmax output for 3 classes
model = tf.keras.Sequential([
    tf.keras.layers.Dense(5, activation='relu', input_shape=(4,)),
    tf.keras.layers.Dense(3, activation='softmax')
])

# Compile the model with loss and optimizer
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Create dummy data: 6 samples, 4 features each
x_train = np.array([[1, 2, 3, 4],
                    [4, 3, 2, 1],
                    [1, 0, 1, 0],
                    [0, 1, 0, 1],
                    [2, 2, 2, 2],
                    [3, 3, 3, 3]], dtype=np.float32)

# Labels for 3 classes (0, 1, or 2)
y_train = np.array([0, 1, 2, 1, 0, 2], dtype=np.int32)

# Train the model for 5 epochs
history = model.fit(x_train, y_train, epochs=5, verbose=0)

# Predict probabilities for a new sample
sample = np.array([[1, 2, 3, 4]], dtype=np.float32)
prediction = model.predict(sample)

print('Predicted probabilities:', prediction)
print('Sum of probabilities:', prediction.sum())

Important Notes

Softmax is best for problems where each input belongs to exactly one class.

The output values are easy to interpret as chances for each class.

Use 'sparse_categorical_crossentropy' loss when labels are integers, not one-hot vectors.

Summary

Softmax turns raw model outputs into probabilities that add to 1.

It is used in the last layer for multi-class classification problems.

The predicted class is the one with the highest probability.