Max Pooling with Argmax

Problem

Max pooling with stride equal to the window size, plus argmax locations.

Slide a k x k window over the image in non-overlapping steps and take the maximum, using the first occurrence on ties. Return [pooled, indices] where indices[i][j] is the [row, col] position of the winning value in the input.

Hint

Track the offset within the window, then add the window origin.

Test cases

case 1

in: [[[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]],2]

expected: [[[6,8],[14,16]],[[[1,1],[1,3]],[[3,1],[3,3]]]]

case 2

in: [[[1,2,3],[4,5,6],[7,8,9]],3]

expected: [[[9]],[[[2,2]]]]

case 3

in: [[[-5,-1],[-3,-2]],2]

expected: [[[-1]],[[[0,1]]]]

case 4

in: [[[2,2],[2,1]],2]

expected: [[[2]],[[[0,0]]]]

Starter codePython

def max_pool_argmax(image, k):
    # Return [pooled, indices]
    # Your code here
    pass

4 test cases · run your code against them in the editor