AdaptiveAvgPool2d

AdaptiveAvgPool2d#

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

Applies a 2D adaptive average pooling over an input signal composed of several input planes.

The output is of size \(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, 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}, 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 square of that size. If a tuple is provided, it specifies (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 5x7
>>> m = AdaptiveAvgPool2d((5, 7))
>>> input = brainstate.random.randn(1, 8, 9, 64)
>>> output = m(input)
>>> output.shape
(1, 5, 7, 64)
>>> # Target output size of 7x7 (square)
>>> m = AdaptiveAvgPool2d(7)
>>> input = brainstate.random.randn(1, 10, 9, 64)
>>> output = m(input)
>>> output.shape
(1, 7, 7, 64)
>>> # Target output size of 10x7
>>> m = AdaptiveAvgPool2d((None, 7))
>>> input = brainstate.random.randn(1, 10, 9, 64)
>>> output = m(input)
>>> output.shape
(1, 10, 7, 64)

AdaptiveAvgPool2d

Contents

AdaptiveAvgPool2d#