Converting Logits to Class Labels

Earlier, a model's outputs were turned into class labels one example at a time. That works for a single example, but a training set or test set contains many examples, and what we usually want is one number that summarizes how often the model is right across all of them. This article builds the pieces that produce that number: first a way to turn each row of model outputs into a single class label, then a reusable function that counts correct predictions over an entire dataset.

A classifier does not usually answer with a class name. For each input it produces a row of numbers, one number per possible class, describing how strongly the model leans toward each class. Before reading the explanation, try the following question. Given the row of scores below, which position would you pick as the model's decision?

The function torch.argmax returns the index of the highest value in a tensor. With dim=1, it scans each row and returns, for every row, the position of that row's largest value. Turning a row of class scores into a single class label is exactly what this does: the returned position is the predicted class.

Notice what argmax does not return. It does not return the value 2.1. It returns the position of that value. The class label is the position, which is why the output of argmax is a set of integers, one per row.

The dimension argument decides the direction of the search. With dim=1, the function scans each row. Changing it 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 using the wrong dimension is one of the most common ways to get label predictions that look plausible but mean something else.

The same argmax operation applies to probabilities and to the raw outputs of the network. 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 argmax can be applied to the logits directly.

Once there are 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. Adding up the True values then counts the correct predictions. The source does this with torch.sum, which counts the True entries in the boolean tensor produced by the comparison.

The source uses a training set of five examples in which all five predictions are correct. In that case the comparison is all True and the count is 5.

Accuracy is the fraction of predictions that were correct. Because the correct count can never exceed the total count of examples, the fraction always falls between 0 and 1. A result of 0 means no predictions matched, and a result of 1 means all of them matched. Expressed as a percentage, the same fraction is multiplied by 100.

The accuracy function does not send the whole dataset through the model at once. A dataset may be far too large for memory, so the function gives the model only one chunk at a time. Each chunk has the same size as the batch size used during training. The function does its work in a single loop over the data loader. On each pass, one batch of features and labels arrives, the model produces logits for that batch, and those logits are converted into predicted labels. The predicted labels are compared with the true labels, and the count of matches is added to a running total of correct predictions. The number of examples in the batch is added to a running count of all examples.

Because the counters carry results forward batch by batch, the final accuracy covers the full dataset even though the model never saw all of it in one call. The same function therefore works for datasets of any size, not just the small ones that fit in memory.

The function is applied to the training data first and then to the test data, using the same model in both calls. Each call prints one value. The two calls differ only in the data loader they receive: the input examples and true labels change, while the function and the model stay the same.

Before you run the two calls, decide what kind of value each one returns and what range that value can take. Because each result is a count of correct predictions divided by a count of all examples, every printed value lies between 0 and 1. A printed value of 0 would mean no predictions matched on that loader, and a printed value of 1 would mean all of them matched. The two printed values can differ because they are computed over different examples.

When an accuracy number surprises you, the mistake almost always sits in one of three places: the dimension passed to argmax, the tensor passed into argmax, or the true labels it is compared against. Work through the pipeline in that order, and check each stage before trusting the final count.