Implement Max Pooling with Indices

Implement max pooling that returns both pooled values and indices for unpooling, a crucial component in CNNs, particularly in encoder-decoder architectures. This operation is essential for preserving spatial information.

Max pooling is a downsampling technique that reduces spatial dimensions by taking the maximum value across each kernel window, defined by $x^2 + y^2 \leq k^2$ , where $k$ is the kernel size. The forward pass involves iterating over the input, applying the max pooling operation to each window.

Here are the key steps:

Divide the input into overlapping or non-overlapping kernel windows.
For each window, find the maximum value and its corresponding index.
Store the index for later use in unpooling operations.

\max_{i, j \in \text{window}} x_{i, j}

This technique is widely used in image segmentation tasks.

Example:

Input:

input: [[[[1, 2], [3, 4]]]]  # 1×1×2×2
kernel_size = 2

Output:

(tensor([[[[4]]]]), tensor([[[[3]]]]))

Reasoning:

2×2 window: [1,2,3,4] Max value: 4 at position (1,1) Flattened index: 1*2 + 1 = 3

Output: max value 4, index 3

Constraints:

input: Tensor (batch, channels, H, W)
kernel_size: Pooling window size
Return: (pooled_output, indices) both same shape as output

Example:

Constraints:

Test Results