Pooling FLOPs

Problem

Count the comparison operations of max pooling: each output compares pool_h * pool_w - 1 candidates, so FLOPs = batch * channels * out_h * out_w * (pool_h * pool_w - 1).

The signature is pooling_flops(batch, channels, out_h, out_w, pool_h, pool_w). Return an integer.

Hint

A window of n elements needs n - 1 comparisons to find the maximum.

Test cases

case 1

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

expected: 36

case 2

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

expected: 3

case 3

in: [2,64,7,7,2,2]

expected: 18816

case 4

in: [1,1,4,4,3,3]

expected: 128

Starter codePython

def pooling_flops(batch, channels, out_h, out_w, pool_h, pool_w):
    # Your code here
    pass

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