Implement Max Pooling with Indices

hardPROCNNsCV: Deep Learning

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, a $k \times k$ square where $k$ is the kernel size (the stride defaults to $k$ , so the windows tile the input without overlapping). 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. The index is the maximum's position in its channel's flattened $H \times W$ input plane, i.e. $\text{row} \cdot W + \text{col}$ (an integer tensor).

\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

Input: x is a $1 \times 1 \times 2 \times 2$ tensor holding $\begin{pmatrix}1 & 2\\ 3 & 4\end{pmatrix}$ with kernel_size $k = 2$ .
Windows: with stride $k$ , each side yields $\lfloor 2/2 \rfloor = 1$ window, so the whole image is one window.
Value: $\max\{1, 2, 3, 4\} = 4$ , kept as a float (4.) because the input is floating point.
Index: the 4 sits at row 1, column 1 of the input plane. Read row by row, its position is $1 \cdot W + 1 = 1 \cdot 2 + 1 = 3$ . Indices are integers, hence 3 with no decimal point.
Shape: both tensors are $1 \times 1 \times 1 \times 1$ ; batch and channel pass through untouched.
Notice the index counts positions in the whole input plane, not inside the window: for a $4 \times 4$ input holding $0, \dots, 15$ , the bottom-right window's maximum 15 is at row 3, column 3, giving index $3 \cdot 4 + 3 = 15$ — exactly where unpooling must put it back.

Constraints

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

Learn this in

Deep LearningCV Study Plan CNNs & Sequence ModelsML Study Plan

More from CV: Deep Learning

Depthwise Separable Convolutionmedium Softmax Cross-Entropy Gradientmedium Batch Normalization Forward Passhard Flatten Feature Mapeasy Max Poolingeasy

Concept primerPREVIEW

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]

PRO

Continue the full explanation

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

PRO

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