Apply Softmax for Predictions

easyPROModelingPyTorch ModelingPytorch

Problem Statement

Convert raw model outputs (logits) into probabilities using Softmax.

Background

Neural networks output raw scores called logits. To interpret these as probabilities, you need to apply Softmax, which converts them to values in [0, 1] that sum to 1 across classes.

Your Task

Write a function apply_softmax(logits) that takes a batch of logits and returns statistics about the resulting probability distribution.

Output Format

Return a dictionary with keys: "probabilities_sum" (list of floats, 1 decimal — should each be ~1.0), "predictions" (list of ints — predicted class per sample), "max_probs" (list of floats, 2 decimals — confidence per sample).

Example

Input

logits = [[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]]

Output

{'probabilities_sum': [1.0, 1.0], 'predictions': [0, 1], 'max_probs': [0.66, 0.8]}

Reasoning

logits is a $2 \times 3$ batch: each row is one sample and its 3 numbers are raw, unbounded class scores, not yet probabilities.
Each row is normalised on its own with $p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$ . Row 1: $e^{2.0} \approx 7.389$ , $e^{1.0} \approx 2.718$ , $e^{0.1} \approx 1.105$ , total $11.213$ , so $p \approx (0.659,\ 0.242,\ 0.099)$ .
Row 2: $e^{0.5} \approx 1.649$ , $e^{2.5} \approx 12.182$ , $e^{0.3} \approx 1.350$ , total $15.181$ , so $p \approx (0.109,\ 0.802,\ 0.089)$ .
probabilities_sum: every row adds to exactly 1 by construction, so both give 1.0 at 1 decimal place.
predictions: the class (counting from 0) holding the largest probability in each row, giving [0, 1].
max_probs: those winning probabilities at 2 decimal places, $0.659 \to 0.66$ and $0.802 \to 0.80$ , which prints as 0.8.
Notice the winner is always the class with the largest logit, because $e^z$ preserves order. Only the confidence depends on the gaps: row 2's lead of $2.0$ yields $0.80$ , row 1's lead of $1.0$ only $0.66$ .

Constraints

Use nn.Softmax(dim=1)
Round probabilities appropriately
Handle batch dimension correctly

Learn this in

PyTorch FundamentalsFoundations

More from Pytorch

Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy Compose Multiple Transformsmedium

Concept primerPREVIEW

1. Background Knowledge

Logits are raw, unnormalized scores output by a neural network's final linear layer, ranging from $[-∞, ∞]$ . They represent each class's relative strength but aren't interpretable as probabilities.

The Softmax function converts logits to a probability distribution:

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}

where $z_i$ is the logit for class $i$ , and $K$ is the number of classes. Key properties:

Outputs in $[0, 1]$
Sums to 1 across classes: $\sum_{i=1}^K$ \text{softmax} $(z_i) = 1$
Higher logits → higher probabilities (exponential scaling amplifies differences)

In PyTorch, nn.Softmax(dim=1) applies this along dimension 1 (class dimension for batch × classes shape), enabling maximum likelihood estimation with cross-entropy loss during training.

Predictions use argmax() on probabilities (or logits, as argmax is invariant to monotonic transforms like softmax).

2. Algorithm Approach

Instantiate Softmax: softmax = nn.Softmax(dim=1)
Apply transformation: probs = softmax(logits) (shape preserved: batch_size × num_classes)
Extract metrics:

Sum: probs.sum(dim=1).tolist() (should be ~1.0 per sample)
Predictions: probs.argmax(dim=1).tolist()
Max probs: probs.max(dim=1).values.tolist()

This is a standard post-processing step in classification pipelines.

3. Step-by-Step Strategy

import torch
import torch.nn as nn

def apply_softmax(logits):
 softmax = nn.Softmax(dim=1)
 probs = softmax(logits)

 probs_sum = probs.sum(dim=1).round(decimals=1).tolist()
 predictions = probs.argmax(dim=1).tolist()
 max_probs = probs.max(dim=1).values.round(decimals=2).tolist()

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.