Max Pooling 2D (MaxPool2D)#

API#

class warp_nn.modules.layers.MaxPool2D(
kernel_size: int | tuple[int, int],
*,
stride: int | tuple[int, int] | None = None,
padding: int | tuple[int, int] = 0,
dilation: int | tuple[int, int] = 1,
ceil_mode: bool = False,
requires_grad: bool = True,
)[source]#

Apply a 2D max pooling.

The output value for each window is the maximum of the input values in the window. The input is implicitly padded with negative infinity values on all sides.

Parameters:

  • kernel_size – The size of the window.

  • stride – The stride of the window. If None, it is equal to kernel_size.

  • padding – The implicit padding added on all sides. It must be at most half of the kernel size.

  • dilation – The spacing between the elements in the window.

  • ceil_mode – Whether to use ceil (instead of floor) to compute the output size. If true, the last windows can extend beyond the (padded) input, as long as they start inside the input or its top/left padding.

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

Raises:

ValueError – If any of the arguments is not valid.

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

Forward pass of the module.

Parameters:

input – The input array, with shape (batch_size, channels, in_height, in_width).

Returns:

The output array, with shape (batch_size, channels, out_height, out_width) where:

\[ \begin{align}\begin{aligned}H_{out} = \left\lfloor \frac{H_{in} + 2 \, \text{padding}[0] - \text{dilation}[0] \, (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} \right\rfloor + 1\\W_{out} = \left\lfloor \frac{W_{in} + 2 \, \text{padding}[1] - \text{dilation}[1] \, (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} \right\rfloor + 1\end{aligned}\end{align} \]

(using ceil instead of floor if ceil_mode is true).

Raises:

ValueError – If the input array’s shape is not supported, or if the output size is too small.