SparKPretrainHead

class mmpretrain.models.heads.SparKPretrainHead(loss, norm_pix=True, patch_size=32)

Pre-training head for SparK.

Parameters:

• loss (dict) – Config of loss.

• norm_pix (bool) – Whether or not normalize target. Defaults to True.

• patch_size (int) – Patch size, equal to downsample ratio of backbone. Defaults to 32.

construct_target(target)

Construct the reconstruction target.

In addition to splitting images into tokens, this module will also normalize the image according to norm_pix. :param target: Image with the shape of B x 3 x H x W :type target: torch.Tensor

Returns:

Tokenized images with the shape of B x L x C

Return type:

torch.Tensor

forward(pred, target, active_mask)

Forward function of MAE head.

Parameters:

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

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

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

Returns:

The reconstruction loss.

Return type:

torch.Tensor

patchify(imgs)

Split images into non-overlapped patches.

Parameters:

imgs (torch.Tensor) – A batch of images, of shape B x C x H x W.

Returns:

Patchified images. The shape is B x L x D.

Return type:

torch.Tensor