torch.nn.Softmax2d
Softmax2d activation function.
Softmax2d applies softmax over the channel dimension for 2D spatial data (images). Given input of shape (N, C, H, W), it applies softmax along the channel dimension (C) at each spatial location (h, w). This is useful for per-pixel channel-wise classification, such as dense prediction tasks where you want probability distributions over classes at each pixel. Common uses include semantic segmentation where each pixel must be classified into one of K classes, or channel-wise attention where channels are weighted as a probability distribution per spatial location.
Core idea: For image tensors (N, C, H, W), apply softmax(x, dim=1) to get probabilities over channels at each pixel. This ensures that for each spatial location, the C channel values sum to 1 and represent a probability distribution. Different from standard Softmax which is usually applied at the end of classification (shape [N, C]), Softmax2d preserves spatial structure.
When to use Softmax2d:
- Semantic segmentation: Dense pixel-wise classification into K classes
- Scene understanding: Dense prediction where each pixel needs a class label
- Per-pixel probability: When you need normalized probabilities per pixel over classes
- Channel-wise attention: Attention weights over channels at each spatial location
- NOT for standard classification: Use standard Softmax for image classification (global)
Input shape requirements:
- 3D input (C, H, W): Applies softmax over C dimension
- 4D input (N, C, H, W): Applies softmax over C dimension (channels)
- Higher dims: Works with any tensor ≥ 3D, applies softmax over dimension 1 for 4D+
Relationship to standard Softmax:
- Softmax(-1): Applied at end of network to class logits, output shape [N, C]
- Softmax2d: Applied to spatial feature maps, output shape [N, C, H, W] (same as input)
- Both produce normalized probabilities, but Softmax2d preserves spatial dimensions
Algorithm: Forward: Apply softmax(x, dim=1) where dim=1 is the channel dimension
- For each spatial location (h, w), softmax over the C channel values
- Result: each pixel has C values that sum to 1 (probability distribution over channels)
- Shape preservation: [N, C, H, W] → [N, C, H, W]
Backward: Chain rule applies softmax gradient at each spatial location
- Gradient flows back through the softmax operation
- Spatial structure preservation: Keeps spatial dimensions unlike standard Softmax.
- Per-pixel probabilities: Each pixel has normalized probability distribution over channels.
- Dense predictions: Standard activation for segmentation and per-pixel tasks.
- Channel dimension: Always operates on dimension 1 (channels) for 4D or higher tensors.
- 3D support: Works with 3D inputs (C, H, W) too, softmax over C dimension.
Examples
// Semantic segmentation: pixel-wise classification
class SemanticSegmentationHead extends torch.nn.Module {
private conv1: torch.nn.Conv2d;
private conv2: torch.nn.Conv2d;
private softmax2d: torch.nn.Softmax2d;
constructor(in_channels: number, num_classes: number) {
super();
this.conv1 = new torch.nn.Conv2d(in_channels, 256, 3, { padding: 1 });
this.conv2 = new torch.nn.Conv2d(256, num_classes, 1);
this.softmax2d = new torch.nn.Softmax2d();
}
forward(x: torch.Tensor): torch.Tensor {
// x: [batch, in_channels, H, W]
x = this.conv1.forward(x); // [batch, 256, H, W]
x = this.conv2.forward(x); // [batch, num_classes, H, W]
return this.softmax2d.forward(x); // [batch, num_classes, H, W] with probabilities per pixel
}
}
// At each pixel: sum of num_classes channel values = 1// Segmentation output: probability per class per pixel
const batch_size = 4, num_classes = 21, height = 256, width = 256;
const logits = torch.randn([batch_size, num_classes, height, width]);
const softmax2d = new torch.nn.Softmax2d();
const probabilities = softmax2d.forward(logits);
// probabilities[0, :, 100, 100] is a [21] vector summing to 1
// Represents P(class_k | pixel (100, 100)) for k=0..20
// Can get predicted class per pixel: argmax(probabilities[0], dim=0)// Comparison: Softmax vs Softmax2d
const batch_logits = torch.randn([32, 10]); // Classification: [batch, classes]
const spatial_logits = torch.randn([32, 10, 64, 64]); // Segmentation: [batch, classes, H, W]
// Classification: softmax over classes
const softmax = new torch.nn.Softmax(-1);
const class_probs = softmax.forward(batch_logits); // [32, 10] - prob distribution per sample
// Segmentation: Softmax2d over channels at each pixel
const softmax2d = new torch.nn.Softmax2d();
const pixel_probs = softmax2d.forward(spatial_logits); // [32, 10, 64, 64] - prob per pixel
// pixel_probs[:, :, h, w].sum(dim=0) == 1 for each (h, w)