RAFT Correlation Volume Size

Problem

Count the entries of a RAFT-style 4D correlation volume pyramid.

Level l correlates two feature maps downsampled by 2^l (each dimension floored, minimum 1), producing (hl * wl) * (hl * wl) pairs. Return the total entry count summed over the given number of levels.

Hint

The single-level all-pairs volume already has (h * w)^2 entries.

Test cases

case 1

in: [32,32,1]

expected: 1048576

case 2

in: [8,8,3]

expected: 4368

case 3

in: [16,8,2]

expected: 17408

case 4

in: [4,4,4]

expected: 274

Starter codePython

def raft_correlation_size(h, w, levels):
    # Your code here
    pass

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