MaxUnpool3d

Contents

MaxUnpool3d#

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

Computes a partial inverse of MaxPool3d.

MaxPool3d is not fully invertible, since the non-maximal values are lost. MaxUnpool3d takes in as input the output of MaxPool3d 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, D_{in}, H_{in}, W_{in}, C)\) or \((D_{in}, H_{in}, W_{in}, C)\)

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

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

    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 = MaxPool3d(2, stride=2, return_indices=True, channel_axis=-1)
>>> unpool = MaxUnpool3d(2, stride=2, channel_axis=-1)
>>> input = brainstate.random.randn(1, 4, 4, 4, 16)
>>> output, indices = pool(input)
>>> unpooled = unpool(output, indices)
>>> # unpooled will have shape (1, 8, 8, 8, 16) with zeros at non-maximal positions