Computing Softmax Probabilities for a Classifier

The same argmax operation works on the probabilities and on the raw outputs of the network, which are called logits.

Programming

A Classifier Output Before Any Decision

A classifier rarely answers with a class name. For each input it produces a row of numbers, one number for every possible class, and those numbers describe how strongly the model leans toward each class. Before anyone can count how many predictions were right, that row has to be reduced to a single decision: the position of its largest value. This article follows that step from the first row of class scores all the way to a final accuracy number.

Try this before reading the explanation. The tensor below has three rows, one per sample, and four class scores in each row. For each row, decide which column position the code will report as the winner. Then check your answer against the reveal.

What do you think happens?

Which position does torch.argmax(probs, dim=1) report for each of the three rows?

tensor([1, 0, 3])
tensor([0, 1, 2])
tensor([0.7, 0.5, 0.5])

Reveal answer

Answer:

tensor([1, 0, 3])

Row 0 has its largest value, 0.7, at column position 1. Row 1 has its largest value, 0.5, at column position 0. Row 2 has its largest value, 0.5, at column position 3, because the final entry in that row is 0.5 and is larger than the other values in that row. The output holds one integer per row, which is the class label for that sample.

Scanning Each Row with argmax

The function torch.argmax returns the index of the highest value. When it is called with dim=1, it scans each row of a two-dimensional tensor and reports the column position of the largest value in that row. Each row contributes exactly one integer to the result, so a tensor with three rows produces three labels.

python

import torch

probs = torch.tensor([[0.1, 0.7, 0.2, 0.0],
                      [0.5, 0.1, 0.3, 0.1],
                      [0.2, 0.2, 0.1, 0.5]])

labels = torch.argmax(probs, dim=1)
print(labels)

import torch

probs = torch.tensor([[0.1, 0.7, 0.2, 0.0],
                      [0.5, 0.1, 0.3, 0.1],
                      [0.2, 0.2, 0.1, 0.5]])

labels = torch.argmax(probs, dim=1)
print(labels)

Output (expected)

tensor([1, 0, 3])

Notes on the code

probs holds three samples with four class scores each
dim=1 tells argmax to scan across each row
the result has one integer per row, not one per value

Trace the scan row by row. In row 0 the values are 0.1, 0.7, 0.2, and 0.0, so the largest value sits at column 1 and that position is reported first. In row 1 the largest value is 0.5 at column 0. In row 2 the largest value is 0.5 at column 3. The three reported positions, 1, 0, and 3, are the predicted class labels for the three samples.

Which column index is picked from each row, and why does the output have one label per row? Each row is scanned on its own, and only its maximum cell contributes a label.

How the dim Argument Sets the Direction

The dim argument decides which direction the search runs. With dim=1 the function scans each row. Changing the argument to dim=0 scans down each column instead, and returns, for every column, the row position of its largest value. That second result answers a different question, so the same tensor can yield a set of numbers that look like labels but mean something else.

python

import torch

probs = torch.tensor([[0.1, 0.7, 0.2, 0.0],
                      [0.5, 0.1, 0.3, 0.1],
                      [0.2, 0.2, 0.1, 0.5]])

print(torch.argmax(probs, dim=0))

import torch

probs = torch.tensor([[0.1, 0.7, 0.2, 0.0],
                      [0.5, 0.1, 0.3, 0.1],
                      [0.2, 0.2, 0.1, 0.5]])

print(torch.argmax(probs, dim=0))

Output (expected)

tensor([1, 0, 1, 2])

Notes on the code

dim=0 scans down each of the four columns
the result has four entries, one per column, not one per sample

Read the output carefully. Column 0 has its largest value, 0.5, in row 1. Column 1 has 0.7 in row 0. Column 2 has 0.3 in row 1. Column 3 has 0.5 in row 2. The result is a valid tensor of row positions, but it has four entries for three samples, so it cannot be a per-sample label list.

Labels Without Softmax

The same argmax operation works on the probabilities and on the raw outputs of the network, which are called logits. Computing softmax first and then calling argmax is a valid path, but it is not required to obtain class labels. The position of the largest value is the same whether you look at the logits or at the probabilities computed from them, so you can apply argmax directly to the logits.

python

import torch

logits = torch.tensor([[ 2.0,  0.5, -1.0,  0.1],
                      [ 0.3,  1.8,  0.2, -0.5],
                      [-0.4,  0.0,  0.9,  2.2]])

probs = torch.softmax(logits, dim=1)

print(torch.argmax(logits, dim=1))
print(torch.argmax(probs, dim=1))

import torch

logits = torch.tensor([[ 2.0,  0.5, -1.0,  0.1],
                      [ 0.3,  1.8,  0.2, -0.5],
                      [-0.4,  0.0,  0.9,  2.2]])

probs = torch.softmax(logits, dim=1)

print(torch.argmax(logits, dim=1))
print(torch.argmax(probs, dim=1))

Output (expected)

tensor([0, 1, 3])
tensor([0, 1, 3])

Notes on the code

logits are the raw network outputs before softmax
probs is computed from logits with softmax across each row
both argmax calls report the same label positions

Before

print(torch.argmax(logits, dim=1))

argmax applied directly to the logits returns tensor([0, 1, 3])

After

print(torch.argmax(torch.softmax(logits, dim=1), dim=1))

argmax applied to the softmax probabilities returns tensor([0, 1, 3])

What happens to the logits at each step, and where does softmax fit in relative to argmax? Softmax is an optional step on the way to the label, and argmax can take the logits directly.

Counting Matches and Computing Accuracy

Once you have predicted labels, the next question is whether they match the true labels. The == operator compares two tensors element by element and returns a tensor of booleans, one for each position, where True means the prediction matched the true label at that position. torch.sum then counts the True entries, and dividing that count by the total number of examples gives the accuracy as a percentage.

python

import torch

predicted = torch.tensor([1, 0, 3])
true_labels = torch.tensor([1, 0, 2])

matches = predicted == true_labels
print(matches)

correct = torch.sum(matches)
print(correct)

accuracy = correct / len(true_labels) * 100
print(accuracy)

import torch

predicted = torch.tensor([1, 0, 3])
true_labels = torch.tensor([1, 0, 2])

matches = predicted == true_labels
print(matches)

correct = torch.sum(matches)
print(correct)

accuracy = correct / len(true_labels) * 100
print(accuracy)

Output (expected)

tensor([ True,  True, False])
tensor(2)
tensor(66.6667)

Notes on the code

== produces one boolean per position
torch.sum counts the True entries
dividing by the number of examples and multiplying by 100 gives a percentage

The accuracy here is 2 correct out of 3 examples, which is about 66.7 percent. The same arithmetic applies to any dataset. In general, accuracy is written as $\frac{\text{correct}}{\text{total}} \times 100$, where correct is the count from torch.sum and total is the number of examples.

Worked example

Accuracy from a Count of Matches

A comparison of predicted labels against true labels has four positions. Three of the four positions are True. What is the accuracy as a percentage?

1
Identify the count of correct predictions
torch.sum over the boolean tensor returns the number of True entries, which is 3.
2
Identify the total number of examples
The comparison has one entry per example, so the total is 4.
3
Divide and convert to a percentage
Compute $\frac{3}{4} \times 100$, which equals 75.

Result

The accuracy is 75 percent.

The source's training set has five examples, and all five predictions are correct. Five correct out of five total gives an accuracy of 100 percent. When every comparison is True, torch.sum returns the full length of the tensor, and dividing by that length gives exactly 100.

How does summing the True and False matches become a count, and how does dividing that count by the total become the accuracy? Each stage passes one value to the next.

Tracing an Unexpected Accuracy

When an accuracy number surprises you, the mistake almost always sits in one of three places: the dim argument passed to argmax, the tensor passed into argmax, or the true labels it is compared against. Work through the pipeline in that order. Print the predicted labels first and check that there is exactly one integer per example. Then print the comparison result and confirm that its length matches the number of examples. Only then should you read the count, because a wrong count is usually a symptom of a wrong label tensor upstream.

MEDIUM

A model returns predicted labels tensor([2, 0, 1, 1, 3]) for five samples, and the true labels are tensor([2, 1, 1, 1, 3]). Work out the boolean comparison tensor, the count from torch.sum, and the accuracy as a percentage. Then state which stage of the pipeline you would check first if the reported accuracy were far lower than you expected.

Hints

Compare the two tensors position by position, and write True or False for each of the five positions.
Count only the True entries, then divide by five and multiply by 100.
If the predicted labels had four entries instead of five, the first stage to check would be the argmax step and its dim argument.

Your attempt (not saved)

Common Mistakes with argmax and Accuracy

Common mistake
Using dim=0 when the rows are samples
What it looks like
torch.argmax(probs, dim=0) on a tensor with three samples and four classes returns four values.
Why it fails
dim=0 scans down each column, so it reports a row position for every class column rather than a class label for every sample.
Do this instead
Use dim=1 so that the search scans each row, and check that the result has one entry per sample.
Common mistake
Assuming softmax must be computed before argmax
What it looks like
Computing torch.softmax on the logits before every call to argmax, just to get labels.
Why it fails
The position of the largest value is the same in the logits and in the probabilities, so the label does not depend on softmax.
Do this instead
Call argmax directly on the logits with dim=1 when only the label is needed, and keep softmax for when you want the probabilities.
Common mistake
Reading the count before checking the labels
What it looks like
Printing the result of torch.sum over the comparison without first checking the predicted label tensor.
Why it fails
A wrong count is usually a symptom of a wrong label tensor produced upstream, so the count alone cannot show where the error started.
Do this instead
Print the predicted labels, then the comparison, then the count, and confirm the length of each matches the number of examples.
Common mistake
Dividing the count by the wrong total, such as the number of classes instead of the number of examples

Key Takeaways

torch.argmax with dim=1 scans each row of a tensor and returns the index of its largest value, so a row of class scores becomes one class label.
dim=0 answers a different question by scanning each column, so the dim argument must match the layout of the samples.
argmax can be applied directly to logits, because softmax does not change which position holds the largest value, so softmax is not required to get labels.
The == operator produces one boolean per example, torch.sum counts the True entries, and dividing by the number of examples gives accuracy as a percentage.
When accuracy looks wrong, check the predicted labels, then the comparison length, then the count, because errors usually start upstream in the label tensor.

Related concepts

More concepts starting with C

A Classifier Output Before Any Decision

Scanning Each Row with argmax

How the dim Argument Sets the Direction

Labels Without Softmax

Counting Matches and Computing Accuracy

Tracing an Unexpected Accuracy

Common Mistakes with argmax and Accuracy

Key Takeaways

Related Topics

Related concepts