Masked Autoencoder Loss

Problem

Compute the masked autoencoder reconstruction loss: mean squared error over the masked patches only, ignoring visible patches.

The signature is masked_autoencoder_loss(pred, target, mask_indices). pred and target are (num_patches, patch_dim) matrices.

Hint

MAE only computes loss on masked patches, which makes training much cheaper.

case 1

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

expected: 2.5

case 2

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

expected: 1

case 3

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

expected: 1

case 4

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

expected: 0

Starter codePython

def masked_autoencoder_loss(pred, target, mask_indices):
    # Your code here
    pass

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