Implement Softmax() Function in NumPy – NumPy Tutorial

In this tutorial, we will use an example to show you how to implement softmax function using numpy. You use code directly.

Softmax function

Softmax function is defined as:

In numpy, if we compute softmax value of an array, we may get underflow and overflow problem. Here is a tutorial:

Implement Softmax Function Without Underflow and Overflow Problem – Deep Learning Tutorial

How to implement softmax function for 1D and 2D array in numpy?

Look at example code:

def softmax(x):
    x_1d = False
    if x.ndim == 1:
        x = np.expand_dims(x, axis = 0)
        x_1d = True
    x = x - np.max(x, axis=1, keepdims=True)
    x = np.exp(x)
    x = x / np.sum(x, axis=1, keepdims=True)
    if x_1d:
        x = np.squeeze(x)
    return x

This function supports 1D and 2D numpy array.

We can test this function as follows:

x = np.array([[1,2,3],[4,5,6]])
print(softmax(x))

x = np.array([1,2,3])
print(softmax(x))

Run this code, you will get:

[[0.09003057 0.24472847 0.66524096]
 [0.09003057 0.24472847 0.66524096]]
[0.09003057 0.24472847 0.66524096]