SimMIM Head Output Dim

Problem

The SimMIM reconstruction head predicts raw pixel values for each masked patch. Given the patch size and output channels, the head's output dimension is patch_size^2 * out_channels.

The signature is simmim_head_output_dim(patch_size, out_channels). Return an integer.

Hint

A linear head over patch tokens predicts all pixels of the patch at once.

Test cases

case 1

in: [16,3]

expected: 768

case 2

in: [4,1]

expected: 16

case 3

in: [32,3]

expected: 3072

case 4

in: [8,4]

expected: 256

Starter codePython

def simmim_head_output_dim(patch_size, out_channels):
    # Your code here
    pass

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