class

Softmax2d

extendsModule

Softmax2d()

source

Softmax activation applied over the channel dimension of a 4-D tensor.

Given a spatial feature map of shape $(N, C, H, W)$ , applies softmax along the channel axis $C$ :

\text{Softmax2d}(x)_{n,c,h,w} = \frac{e^{x_{n,c,h,w}}}{\sum_{c'} e^{x_{n,c',h,w}}}

Equivalent to Softmax(dim=-3) — channels are the third-from-last axis. Used in dense prediction heads (semantic segmentation, saliency maps) where each spatial location's channel scores form a categorical distribution.

Notes

Input: $(N, C, H, W)$ — must be exactly 4-D.
Output: $(N, C, H, W)$ — values along the channel dimension sum to 1 at every spatial location.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> m = nn.Softmax2d()
>>> x = lucid.randn(1, 5, 4, 4)
>>> out = m(x)
>>> out.shape
(1, 5, 4, 4)
>>> # Verify that channel probabilities sum to 1 at each pixel
>>> import lucid
>>> x = lucid.randn(2, 10, 8, 8)
>>> out = nn.Softmax2d()(x)
>>> # out.sum(dim=1) should be all-ones spatially
>>> out.sum(dim=1).shape
(2, 8, 8)

Methods (1)

forward

→Tensor

forward(x: Tensor)

source

Apply the activation function element-wise.

Parameters

inputTensor

Input tensor of arbitrary shape.

Returns

Tensor

Output tensor of the same shape as input.