Skip to content

MaxPool2d

simplegrad.nn.pooling.MaxPool2d

Bases: Module

2D max pooling layer.

Slides a window over the input and keeps the maximum value in each window.

Parameters:

Name Type Description Default
kernel_size int | tuple

Pooling window size. Int or (kH, kW).

required
stride int | tuple

Step between windows. Int or (sH, sW).

required
pad_width int | tuple[int]

Padding. Int (all sides) or (top, bottom, left, right). Defaults to 0.

0
pad_mode str

Padding mode. Defaults to "constant".

'constant'
pad_value int

Fill value for constant padding. Defaults to 0.

0
Source code in simplegrad/nn/pooling.py
class MaxPool2d(Module):
    """2D max pooling layer.

    Slides a window over the input and keeps the maximum value in each window.

    Args:
        kernel_size: Pooling window size. Int or ``(kH, kW)``.
        stride: Step between windows. Int or ``(sH, sW)``.
        pad_width: Padding. Int (all sides) or ``(top, bottom, left, right)``.
            Defaults to 0.
        pad_mode: Padding mode. Defaults to ``"constant"``.
        pad_value: Fill value for constant padding. Defaults to 0.
    """

    def __init__(
        self,
        kernel_size: int | tuple,
        stride: int | tuple,
        # 2 options for padding:
        # int - same dimention for all directions
        # tuple - (top, bottom, left, right)
        pad_width: int | tuple[int] = 0,
        pad_mode: str = "constant",
        pad_value: int = 0,
    ):
        super().__init__()
        assert ((isinstance(kernel_size, int)) and kernel_size > 0) or (
            (isinstance(kernel_size, tuple) and all(k > 0 for k in kernel_size))
            and len(kernel_size) == 2
        ), "kernel_size must be a positive integer or a tuple of positive integers"
        self.kernel_size = (
            kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        )
        assert ((isinstance(stride, int)) and stride > 0) or (
            (isinstance(stride, tuple) and all(s > 0 for s in stride)) and len(stride) == 2
        ), "stride must be a positive integer or a tuple of positive integers"
        self.stride = stride if isinstance(stride, tuple) else (stride, stride)

        self.pad_width = pad_width
        self.pad_mode = pad_mode
        self.pad_value = pad_value

    def forward(self, x: Tensor) -> Tensor:
        """Apply max pooling to the input.

        Args:
            x: Input of shape ``(batch, channels, H, W)`` or ``(channels, H, W)``.

        Returns:
            Pooled output of shape ``(batch, channels, out_H, out_W)``.
        """
        return max_pool2d(
            x,
            kernel_size=self.kernel_size,
            stride=self.stride,
            pad_width=self.pad_width,
            pad_mode=self.pad_mode,
            pad_value=self.pad_value,
        )

forward(x: Tensor) -> Tensor

Apply max pooling to the input.

Parameters:

Name Type Description Default
x Tensor

Input of shape (batch, channels, H, W) or (channels, H, W).

required

Returns:

Type Description
Tensor

Pooled output of shape (batch, channels, out_H, out_W).

Source code in simplegrad/nn/pooling.py
def forward(self, x: Tensor) -> Tensor:
    """Apply max pooling to the input.

    Args:
        x: Input of shape ``(batch, channels, H, W)`` or ``(channels, H, W)``.

    Returns:
        Pooled output of shape ``(batch, channels, out_H, out_W)``.
    """
    return max_pool2d(
        x,
        kernel_size=self.kernel_size,
        stride=self.stride,
        pad_width=self.pad_width,
        pad_mode=self.pad_mode,
        pad_value=self.pad_value,
    )