MaxUnpool2d

Contents

MaxUnpool2d#

class brainstate.nn.MaxUnpool2d(kernel_size, stride=None, padding=0, channel_axis=-1, name=None, in_size=None)#

Computes a partial inverse of MaxPool2d.

MaxPool2d is not fully invertible, since the non-maximal values are lost. MaxUnpool2d takes in as input the output of MaxPool2d including the indices of the maximal values and computes a partial inverse in which all non-maximal values are set to zero.

Shape:

  • Input: \((N, H_{in}, W_{in}, C)\) or \((H_{in}, W_{in}, C)\)

  • Output: \((N, H_{out}, W_{out}, C)\) or \((H_{out}, W_{out}, C)\), where

    \[H_{out} = (H_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{kernel\_size}[0]\]
    \[W_{out} = (W_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] + \text{kernel\_size}[1]\]

    or as given by output_size in the call operator

Parameters:

  • kernel_size (int | Sequence[int] | integer | Sequence[integer]) – Size of the max pooling window.

  • stride (int | Sequence[int]) – Stride of the max pooling window. Default: kernel_size

  • padding (int | Tuple[int, ...]) – Padding that was added to the input. Default: 0

  • channel_axis (int | None) – Axis of the channels. Default: -1

  • name (str | None) – Name of the module.

  • in_size (int | Sequence[int] | integer | Sequence[integer] | None) – Input size for shape inference.

Examples

>>> import brainstate
>>> # Create pooling and unpooling layers
>>> pool = MaxPool2d(2, stride=2, return_indices=True, channel_axis=-1)
>>> unpool = MaxUnpool2d(2, stride=2, channel_axis=-1)
>>> input = brainstate.random.randn(1, 4, 4, 16)
>>> output, indices = pool(input)
>>> unpooled = unpool(output, indices)
>>> # unpooled will have shape (1, 8, 8, 16) with zeros at non-maximal positions
Contents