PyTorch Softmax and Probabilities

Softmax is a function that turns raw scores into probabilities. The probabilities it makes always add up to 1.

One row of scores

Softmax turns the three scores into three probabilities.

import torch

scores = torch.tensor([2.0, 1.0, 0.0])
print(torch.softmax(scores, dim=0))

Every row sums to 1

Rows of a table are scored one row at a time with dim=1.

import torch

scores = torch.tensor([[2.0, 1.0, 0.0],
                       [0.0, 1.0, 3.0],
                       [0.0, 2.0, 1.0]])
probabilities = torch.softmax(scores, dim=1)
print(probabilities.sum(dim=1))

Reading the prediction off

Argmax finds the position of the biggest probability. That position is the predicted class.

import torch

scores = torch.tensor([0.0, 1.0, 3.0])
probabilities = torch.softmax(scores, dim=0)
print(torch.argmax(probabilities))

Example

import torch

scores = torch.tensor([[2.0, 1.0, 0.0],
                       [0.0, 1.0, 3.0],
                       [0.0, 2.0, 1.0]])
probabilities = torch.softmax(scores, dim=1)
predicted = torch.argmax(probabilities, dim=1)

print("Probabilities:")
print(probabilities)
print("Row sums:")
print(probabilities.sum(dim=1))
print(f"Predicted classes: {predicted.tolist()}")

Key points

Probability: a number between 0 and 1 that says how likely a choice is
Use torch.softmax(scores, dim=1) to turn each row of scores into probabilities
The dim argument picks the direction, and dim=1 runs across each row
The biggest score becomes the biggest probability
Use torch.argmax to find the position of the biggest probability
Softmax keeps the order of the scores and rescales them into shares of 1

In this example, three rows of scores go through softmax with dim=1. Each row becomes a row of probabilities. Every row of probabilities sums to 1, as the printed sums show. Argmax then reads the biggest position in each row. The predicted classes print as 0, 2, and 1.

Note

Softmax never changes which score is biggest, so argmax gives the same answer before and after

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

scores = torch.tensor([[2.0, 1.0, 0.0],
                       [0.0, 1.0, 3.0],
                       [0.0, 2.0, 1.0]])
probabilities = torch.softmax(scores, dim=1)
predicted = torch.argmax(probabilities, dim=1)

print("Probabilities:")
print(probabilities)
print("Row sums:")
print(probabilities.sum(dim=1))
print(f"Predicted classes: {predicted.tolist()}")

Exercise

What do the probabilities in one softmax row always add up to?

1 100 The biggest score 0