Global Average Pooling (GlobalAvgPool)#

API#

class warp_nn.modules.layers.GlobalAvgPool(*, requires_grad: bool = True)[source]#

Bases: Module

Apply a global average pooling over the spatial dimensions of the input.

The output value for each sample and channel is the average of all the input values in the spatial dimensions.

Parameters:

requires_grad – Whether the cached output arrays of the module require gradients.

__call__(
input: array,
) → array[source]#

Forward pass of the module.

Parameters:

input – The input array, with shape (batch_size, channels, *spatial), with 1 or 2 spatial dimensions.

Returns:

The output array, with shape (batch_size, channels, 1) or (batch_size, channels, 1, 1).

Raises:

ValueError – If the input array’s shape is not supported.