Softmax2d

Contents

Softmax2d#

class brainstate.nn.Softmax2d(name=None)#

Applies SoftMax over features to each spatial location.

When given an image of Channels x Height x Width, it will apply Softmax to each location \((Channels, h_i, w_j)\)

Shape#

  • Input: \((N, C, H, W)\) or \((C, H, W)\).

  • Output: \((N, C, H, W)\) or \((C, H, W)\) (same shape as input)

returns:

A Tensor of the same dimension and shape as the input with values in the range [0, 1]

rtype:

Tensor

Examples

>>> import brainstate.nn as nn
>>> import brainstate
>>> m = nn.Softmax2d()
>>> # you softmax over the 2nd dimension
>>> x = brainstate.random.randn(2, 3, 12, 13)
>>> output = m(x)
Contents