PixelReconstructionLoss

class mmpretrain.models.losses.PixelReconstructionLoss(criterion, channel=None)

Loss for the reconstruction of pixel in Masked Image Modeling.

This module measures the distance between the target image and the reconstructed image and compute the loss to optimize the model. Currently, This module only provides L1 and L2 loss to penalize the reconstructed error. In addition, a mask can be passed in the forward function to only apply loss on visible region, like that in MAE.

Parameters:

• criterion (str) – The loss the penalize the reconstructed error. Currently, only supports L1 and L2 loss

• channel (int, optional) – The number of channels to average the reconstruction loss. If not None, the reconstruction loss will be divided by the channel. Defaults to None.

forward(pred, target, mask=None)

Forward function to compute the reconstrction loss.

Parameters:

• pred (torch.Tensor) – The reconstructed image.

• target (torch.Tensor) – The target image.

• mask (torch.Tensor) – The mask of the target image.

Returns:

The reconstruction loss.

Return type:

torch.Tensor