Activation Memory Size

Problem

Compute the memory footprint of an activation tensor in bytes.

Return batch * channels * h * w * bytes_per_value, where bytes_per_value is 4 for float32 and 2 for float16.

case 1

in: [1,64,56,56,4]

expected: 802816

case 2

in: [8,3,224,224,4]

expected: 4816896

case 3

in: [2,16,7,7,2]

expected: 3136

Starter codePython

def activation_memory(batch, channels, h, w, bytes_per_value):
    # Your code here
    pass

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