AdaptiveMaxPool3d

AdaptiveMaxPool3d#

class brainstate.nn.AdaptiveMaxPool3d(target_size, channel_axis=-1, name=None, in_size=None)#

Applies a 3D adaptive max pooling over an input signal composed of several input planes.

The output is of size \(D_{out} \times H_{out} \times W_{out}\), for any input size. The number of output features is equal to the number of input planes.

Adaptive pooling automatically computes the kernel size and stride to achieve the desired output size, making it useful for creating fixed-size representations from variable-sized inputs.

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}, H_{out}, W_{out})=\text{target\_size}\).

Parameters:
  • target_size (int | Sequence[int]) – The target output size. If a single integer is provided, the output will be a cube of that size. If a tuple is provided, it specifies (D_out, H_out, W_out). Use None for dimensions that should not be pooled.

  • channel_axis (int | None) – Axis of the spatial channels for which pooling is skipped. If None, there is no channel axis. Default: -1

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

  • in_size (Sequence[int] | None) – The shape of the input tensor for shape inference.

Examples

>>> import brainstate
>>> # Target output size of 5x7x9
>>> m = AdaptiveMaxPool3d((5, 7, 9))
>>> input = brainstate.random.randn(1, 8, 9, 10, 64)
>>> output = m(input)
>>> output.shape
(1, 5, 7, 9, 64)
>>> # Target output size of 7x7x7 (cube)
>>> m = AdaptiveMaxPool3d(7)
>>> input = brainstate.random.randn(1, 10, 9, 8, 64)
>>> output = m(input)
>>> output.shape
(1, 7, 7, 7, 64)
>>> # Target output size of 7x9x8
>>> m = AdaptiveMaxPool3d((7, None, None))
>>> input = brainstate.random.randn(1, 10, 9, 8, 64)
>>> output = m(input)
>>> output.shape
(1, 7, 9, 8, 64)

See also

MaxPool3d

Non-adaptive 3D max pooling.

AdaptiveAvgPool3d

Adaptive 3D average pooling.