AdaptiveAvgPool3d

AdaptiveAvgPool3d#

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

Applies a 3D adaptive average 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 = AdaptiveAvgPool3d((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 = AdaptiveAvgPool3d(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 = AdaptiveAvgPool3d((7, None, None))
>>> input = brainstate.random.randn(1, 10, 9, 8, 64)
>>> output = m(input)
>>> output.shape
(1, 7, 9, 8, 64)

See also

AvgPool3d

Non-adaptive 3D average pooling.

AdaptiveMaxPool3d

Adaptive 3D max pooling.