MAEPretrainHead

class mmpretrain.models.heads.MAEPretrainHead(loss, norm_pix=False, patch_size=16)

Head for MAE Pre-training.

参数:

• loss (dict) – Config of loss.

• norm_pix_loss (bool) – Whether or not normalize target. Defaults to False.

• patch_size (int) – Patch size. Defaults to 16.

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.

参数:

target (torch.Tensor) – Image with the shape of B x 3 x H x W

返回:

Tokenized images with the shape of B x L x C

返回类型:

torch.Tensor

loss(pred, target, mask)

Generate loss.

参数:

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

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

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

返回:

The reconstruction loss.

返回类型:

torch.Tensor

patchify(imgs)

Split images into non-overlapped patches.

参数:

imgs (torch.Tensor) – A batch of images. The shape should be \((B, 3, H, W)\).

返回:

Patchified images. The shape is \((B, L, \text{patch_size}^2 \times 3)\).

返回类型:

torch.Tensor

unpatchify(x)

Combine non-overlapped patches into images.

参数:

x (torch.Tensor) – The shape is \((B, L, \text{patch_size}^2 \times 3)\).

返回:

The shape is \((B, 3, H, W)\).

返回类型:

torch.Tensor