📘

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

More from CV: Deep Learning

Depthwise Separable ConvolutionMedium Softmax Cross-Entropy GradientMedium Batch Normalization Forward PassHard Flatten Feature MapEasy Max PoolingEasy

You want a max pooling layer that returns both the pooled values and, for each output, the index of the winning element inside its pooling window. This is the variant used in encoder–decoder CNNs (e.g., SegNet) so that a later unpooling layer can place values back at the right spatial positions.

1. Background Knowledge (key concepts)

Max pooling in CNNs takes an input feature map (typically shape: batch × channels × height × width) and, for each small spatial window (e.g., 2×2), outputs the maximum value in that window. This reduces spatial resolution and keeps the strongest activations, acting as a form of downsampling + invariance. Formally, for a window $W$ , max pooling computes

y = \max_{(i,j) \in W} x_{ij}.

In encoder–decoder architectures for dense prediction (segmentation), we sometimes want to invert this downsampling approximately. A max-unpooling layer upsamples by putting each pooled value back at the position of its original maximum and filling other positions with zeros. To do this correctly, the forward max pooling must remember the argmax position inside each window. These positions are often stored as flattened indices from $0$ to $(k_h \cdot k_w - 1)$ , where $k_h, k_w$ are the kernel (window) dimensions.

2. Algorithm / General Approach

General pattern for max pooling with indices:

Slide a fixed-size window over the input with given stride and (optionally) padding.
For each window:

Scan all elements in the window.
Track:
the maximum value encountered so far,
the position (row, col) of this maximum within the window.

Write:

the maximum value to the corresponding output location,
the flattened index of the winning (row, col) to a parallel indices tensor.

Flattening a window position $(r, c)$ (0-based within the window) is usually:

\text{index} = r \cdot k_w + c

for a $k_h \times k_w$ window.

3. Step-by-Step Implementation Strategy

Assume input shape: N x C x H x W, kernel (kH, kW), stride (sH, sW), no padding for simplicity.

Compute output shape

$H_{out} = \lfloor$ \frac{H - kH}{sH} $\rfloor + 1$
$W_{out} = \lfloor$ \frac{W - kW}{sW} $\rfloor + 1$

Allocate outputs

output: shape N x C x H_out x W_out
indices: same shape, storing int (flattened index per output position)

Loop over batch and channels

for n in range(N):
for c in range(C):
# process 2D map input[n, c]

👑

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.