import torch
import lightning as L
from torch import nn, Tensor
from typing import Optional
from ..nets.image.vit_local import VisionTransformer, Block
from ...utils.position_embedding import get_2d_sincos_pos_embed
from einops import rearrange
[docs]
class MaskedAutoEncoderViT(L.LightningModule):
def __init__(
self,
backbone: Optional[VisionTransformer] = None,
decoder_embed_dim=512,
decoder_depth=8,
decoder_num_heads=16,
mlp_ratio=4,
norm_pix_loss=False,
lr: float = 1e-4,
weight_decay: float = 0,
):
super().__init__()
self.lr = lr
self.weight_decay = weight_decay
self.backbone = backbone or VisionTransformer()
embed_dim = self.backbone.embed_dim
patch_size = self.backbone.patch_size
in_channels = self.backbone.in_channels
num_patches = self.backbone.patch_embed.num_patches
self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim)
self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
self.decoder_pos_embed = nn.Parameter(
torch.zeros(1, num_patches + 1, decoder_embed_dim),
requires_grad=False,
)
self.decoder_blocks = nn.ModuleList(
[
Block(
decoder_embed_dim,
decoder_num_heads,
mlp_ratio,
qkv_bias=True,
norm_layer=nn.LayerNorm,
)
for _ in range(decoder_depth)
]
)
self.decoder_norm = nn.LayerNorm(decoder_embed_dim)
self.decoder_pred = nn.Linear(
decoder_embed_dim, patch_size[0] * patch_size[1] * in_channels
)
self.norm_pix_loss = norm_pix_loss
self.initialize_weights()
[docs]
def initialize_weights(self):
pos_embed = get_2d_sincos_pos_embed(
self.backbone.pos_embed.shape[-1],
self.backbone.patch_embed.grid_size,
True,
)
self.backbone.pos_embed.data.copy_(
torch.from_numpy(pos_embed).float().unsqueeze(0)
)
self.backbone.pos_embed.requires_grad_(False)
decoder_pos_embed = get_2d_sincos_pos_embed(
self.decoder_pos_embed.shape[-1],
self.backbone.patch_embed.grid_size,
True,
)
self.decoder_pos_embed.data.copy_(
torch.from_numpy(decoder_pos_embed).float().unsqueeze(0)
)
w = self.backbone.patch_embed.proj.weight.data
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
nn.init.normal_(self.backbone.cls_token, std=0.02)
nn.init.normal_(self.mask_token, std=0.02)
self.apply(self._init_weights)
[docs]
def _init_weights(self, m):
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1)
[docs]
def patchify(self, imgs: Tensor) -> Tensor:
h, w = self.backbone.patch_embed.patch_size
*_, H, W = imgs.shape
assert H % h == 0 and W % w == 0
return rearrange(imgs, "B C (Hp h) (Wp w) -> B (Hp Wp) (h w C)", h=h, w=w)
[docs]
def unpatchify(self, x: Tensor) -> Tensor:
h, w = self.backbone.patch_embed.patch_size
Hp, Wp = self.backbone.patch_embed.grid_size
L = x.shape[1]
assert L == Hp * Wp
return rearrange(
x, "B (Hp Wp) (h w C) -> B C (Hp h) (Wp w)", h=h, w=w, Hp=Hp, Wp=Wp
)
[docs]
def random_masking(self, x, mask_ratio):
"""
Perform per-sample random masking by per-sample shuffling.
Per-sample shuffling is done by argsort random noise.
x: [N, L, D], sequence
"""
N, L, D = x.shape # batch, length, dim
len_keep = int(L * (1 - mask_ratio))
noise = torch.rand(N, L, device=x.device) # noise in [0, 1]
# sort noise for each sample
ids_shuffle = torch.argsort(
noise, dim=1
) # ascend: small is keep, large is remove
ids_restore = torch.argsort(ids_shuffle, dim=1)
# keep the first subset
ids_keep = ids_shuffle[:, :len_keep]
x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
# generate the binary mask: 0 is keep, 1 is remove
mask = torch.ones([N, L], device=x.device)
mask[:, :len_keep] = 0
# unshuffle to get the binary mask
mask = torch.gather(mask, dim=1, index=ids_restore)
return x_masked, mask, ids_restore
[docs]
def forward_encoder(self, x, mask_ratio):
# embed patches
x = self.backbone.patch_embed(x)
# add pos embed w/o cls token
x = x + self.backbone.pos_embed[:, 1:, :]
# masking: length -> length * mask_ratio
x, mask, ids_restore = self.random_masking(x, mask_ratio)
# append cls token
cls_token = self.backbone.cls_token + self.backbone.pos_embed[:, :1, :]
cls_tokens = cls_token.expand(x.shape[0], -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
# apply Transformer blocks
x = self.backbone.blocks(x)
return x, mask, ids_restore
[docs]
def forward_decoder(self, x, ids_restore):
# embed tokens
x = self.decoder_embed(x)
# append mask tokens to sequence
mask_tokens = self.mask_token.repeat(
x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1
)
x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1) # no cls token
x_ = torch.gather(
x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2])
) # unshuffle
x = torch.cat([x[:, :1, :], x_], dim=1) # append cls token
# add pos embed
x = x + self.decoder_pos_embed
# apply Transformer blocks
for blk in self.decoder_blocks:
x = blk(x)
x = self.decoder_norm(x)
# predictor projection
x = self.decoder_pred(x)
# remove cls token
x = x[:, 1:, :]
return x
[docs]
def forward_loss(self, imgs, pred, mask):
"""
imgs: [N, 3, H, W]
pred: [N, L, p*p*3]
mask: [N, L], 0 is keep, 1 is remove,
"""
target = self.patchify(imgs)
if self.norm_pix_loss:
mean = target.mean(dim=-1, keepdim=True)
var = target.var(dim=-1, keepdim=True)
target = (target - mean) / (var + 1.0e-6) ** 0.5
loss = (pred - target) ** 2
loss = loss.mean(dim=-1) # [N, L], mean loss per patch
loss = (loss * mask).sum() / mask.sum() # mean loss on removed patches
return loss
[docs]
def forward(self, imgs, mask_ratio=0.75):
latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
pred = self.forward_decoder(latent, ids_restore) # [N, L, p*p*3]
loss = self.forward_loss(imgs, pred, mask)
return loss, pred, mask
[docs]
def _single_step(self, batch: torch.Tensor, step_name: str) -> torch.Tensor:
loss, _, _ = self.forward(batch)
self.log(f"{step_name}_loss", loss)
return loss
[docs]
def training_step(self, batch: torch.Tensor, batch_idx: int):
return self._single_step(batch, "train")
[docs]
def validation_step(self, batch: torch.Tensor, batch_idx: int):
return self._single_step(batch, "val")
[docs]
def test_step(self, batch: torch.Tensor, batch_idx: int):
return self._single_step(batch, "test")