Turn class logits into normalized class probabilities
Classification
PyTorch · Function
F.softmax
Applies softmax along a selected tensor dimension, converting each slice into values from 0 to 1 that sum to 1. An optional dtype controls casting before computation.
How it works
F.softmax applies the exponential normalization independently to each slice along the chosen dimension. It returns values from 0 to 1 whose sum along that dimension is 1; an optional dtype casts the input before computation, which can help avoid overflow.
Syntax
torch.nn.functional.softmax(input, dim=None, _stacklevel=3, dtype=None)Parameters
| Name | Type | Description |
|---|---|---|
input | Tensor | Input tensor. |
dim | int | Dimension along which softmax is computed. |
dtype | torch.dtype, optional | Optional output dtype; the input is cast to this type before computation. |
Returns
Tensor
Key points
- Normalization occurs independently for every slice along dim, producing values between 0 and 1 that sum to 1.
- The dim argument chooses the axis used for normalization.
- dtype controls the desired output data type.
- When dtype is supplied, casting happens before softmax computation.
- Supplying an appropriate dtype can reduce the risk of overflow during computation.
Common mistakes
- Do not pass softmax output directly to NLLLoss; use log_softmax instead.
- Prefer log_softmax when a downstream loss needs log-probabilities.
Example: Convert logits to probabilities
import torch
import torch.nn.functional as F
logits = torch.tensor([[1.0, 2.0, 0.5], [0.0, 0.0, 1.0]])
probabilities = F.softmax(logits, dim=1)
print(probabilities)
print(probabilities.sum(dim=1))Converts each row of logits into probabilities whose row sums are 1.
When to use it
Convert scores along a chosen axis into normalized weights
Attention and weighting
Compute softmax using a selected dtype when overflow prevention is needed
Numerical preparation
Interview questions
QHow does F.softmax transform an input tensor?
It exponentiates each value and divides by the sum of exponentials across the selected dimension. The resulting slice contains values in [0, 1] and sums to 1.
Follow-up: What does dim control? It selects the dimension whose slices are normalized independently; choosing the wrong dimension produces probabilities over the wrong axis.
QWhat is the purpose of dim?
dim identifies the axis along which softmax is computed. For a batch of class logits shaped [batch, classes], dim=1 normalizes each example's class scores.
Follow-up: How would you normalize the final dimension without hard-coding its index? Use dim=-1 when that dimension is the intended class or feature axis.
QWhy might dtype or log_softmax matter?
dtype optionally specifies the returned tensor's data type, and the input is cast to it before the operation. The documentation notes this can help prevent data-type overflow.
Follow-up: When should log_softmax be used instead? Use it when log-probabilities are required, especially before NLLLoss, because it is faster and has better numerical properties.
Related
torch.nn.Softmaxtorch.nn.functional.log_softmaxtorch.softmaxtorch.nn.NLLLoss