# --------------------------------------------------------
# References:
# MAE: https://github.com/IcarusWizard/MAE
# --------------------------------------------------------
import torch
import timm
import numpy as np
from einops import repeat, rearrange
from einops.layers.torch import Rearrange
from timm.models.layers import trunc_normal_
from timm.models.vision_transformer import Block
def random_indexes(size : int):
forward_indexes = np.arange(size)
np.random.shuffle(forward_indexes)
backward_indexes = np.argsort(forward_indexes)
return forward_indexes, backward_indexes
def take_indexes(sequences, indexes):
return torch.gather(sequences, 0, repeat(indexes, 't b -> t b c', c=sequences.shape[-1]))
class PatchShuffle(torch.nn.Module):
def __init__(self, ratio) -> None:
super().__init__()
self.ratio = ratio
def forward(self, patches : torch.Tensor):
T, B, C = patches.shape
remain_T = int(T * (1 - self.ratio))
indexes = [random_indexes(T) for _ in range(B)]
forward_indexes = torch.as_tensor(np.stack([i[0] for i in indexes], axis=-1), dtype=torch.long).to(patches.device)
backward_indexes = torch.as_tensor(np.stack([i[1] for i in indexes], axis=-1), dtype=torch.long).to(patches.device)
patches = take_indexes(patches, forward_indexes)
patches = patches[:remain_T]
return patches, forward_indexes, backward_indexes
class MAE_Encoder(torch.nn.Module):
def __init__(self,
image_size=32,
patch_size=2,
emb_dim=192,
num_layer=12,
num_head=3,
mask_ratio=0.75,
) -> None:
super().__init__()
self.cls_token = torch.nn.Parameter(torch.zeros(1, 1, emb_dim))
self.pos_embedding = torch.nn.Parameter(torch.zeros((image_size // patch_size) ** 2, 1, emb_dim))
self.shuffle = PatchShuffle(mask_ratio)
self.patchify = torch.nn.Conv2d(3, emb_dim, patch_size, patch_size)
self.transformer = torch.nn.Sequential(*[Block(emb_dim, num_head) for _ in range(num_layer)])
self.layer_norm = torch.nn.LayerNorm(emb_dim)
self.init_weight()
def init_weight(self):
trunc_normal_(self.cls_token, std=.02)
trunc_normal_(self.pos_embedding, std=.02)
def forward(self, img):
patches = self.patchify(img)
patches = rearrange(patches, 'b c h w -> (h w) b c')
patches = patches + self.pos_embedding
patches, forward_indexes, backward_indexes = self.shuffle(patches)
patches = torch.cat([self.cls_token.expand(-1, patches.shape[1], -1), patches], dim=0)
patches = rearrange(patches, 't b c -> b t c')
features = self.layer_norm(self.transformer(patches))
features = rearrange(features, 'b t c -> t b c')
return features, backward_indexes
class MAE_Decoder(torch.nn.Module):
def __init__(self,
image_size=32,
patch_size=2,
emb_dim=192,
num_layer=4,
num_head=3,
) -> None:
super().__init__()
self.mask_token = torch.nn.Parameter(torch.zeros(1, 1, emb_dim))
self.pos_embedding = torch.nn.Parameter(torch.zeros((image_size // patch_size) ** 2 + 1, 1, emb_dim))
self.transformer = torch.nn.Sequential(*[Block(emb_dim, num_head) for _ in range(num_layer)])
self.head = torch.nn.Linear(emb_dim, 3 * patch_size ** 2)
self.patch2img = Rearrange('(h w) b (c p1 p2) -> b c (h p1) (w p2)', p1=patch_size, p2=patch_size, h=image_size//patch_size)
self.init_weight()
def init_weight(self):
trunc_normal_(self.mask_token, std=.02)
trunc_normal_(self.pos_embedding, std=.02)
def forward(self, features, backward_indexes):
T = features.shape[0]
backward_indexes = torch.cat([torch.zeros(1, backward_indexes.shape[1]).to(backward_indexes), backward_indexes + 1], dim=0)
features = torch.cat([features, self.mask_token.expand(backward_indexes.shape[0] - features.shape[0], features.shape[1], -1)], dim=0)
features = take_indexes(features, backward_indexes)
features = features + self.pos_embedding
features = rearrange(features, 't b c -> b t c')
features = self.transformer(features)
features = rearrange(features, 'b t c -> t b c')
features = features[1:] # remove global feature
patches = self.head(features)
mask = torch.zeros_like(patches)
mask[T-1:] = 1
mask = take_indexes(mask, backward_indexes[1:] - 1)
img = self.patch2img(patches)
mask = self.patch2img(mask)
return img, mask
class MAE_ViT(torch.nn.Module):
def __init__(self,
image_size=32,
patch_size=2,
emb_dim=192,
encoder_layer=12,
encoder_head=3,
decoder_layer=4,
decoder_head=3,
mask_ratio=0.75,
) -> None:
super().__init__()
self.encoder = MAE_Encoder(image_size, patch_size, emb_dim, encoder_layer, encoder_head, mask_ratio)
self.decoder = MAE_Decoder(image_size, patch_size, emb_dim, decoder_layer, decoder_head)
def forward(self, img):
features, backward_indexes = self.encoder(img)
predicted_img, mask = self.decoder(features, backward_indexes)
return predicted_img, mask
class ViT_Classifier(torch.nn.Module):
'''
A simple image classification task acts as a head for ViT, allowing fine-tuning on downstream tasks.
We didn't directly use the MAE_ViT encoder because we need to add a classification head.
The Masked Autoencoder uses only some patches as input, which means it lacks the global information of the image,
making it unsuitable for classification.
'''
def __init__(self, encoder : MAE_Encoder, dropout_p, num_classes=10) -> None:
super().__init__()
self.dropout_p = dropout_p
self.cls_token = encoder.cls_token
self.pos_embedding = encoder.pos_embedding
self.patchify = encoder.patchify
self.transformer = encoder.transformer
self.layer_norm = encoder.layer_norm
self.dropout = torch.nn.Dropout(dropout_p) # Add dropout layer
self.head = torch.nn.Linear(self.pos_embedding.shape[-1], num_classes)
def forward(self, img):
patches = self.patchify(img)
patches = rearrange(patches, 'b c h w -> (h w) b c')
patches = patches + self.pos_embedding
patches = torch.cat([self.cls_token.expand(-1, patches.shape[1], -1), patches], dim=0)
patches = rearrange(patches, 't b c -> b t c')
features = self.layer_norm(self.transformer(patches))
# t is the number of patches, b is the batch size, c is the number of features
features = rearrange(features, 'b t c -> t b c')
if self.dropout_p > 0:
features = self.dropout(features) # Apply dropout before the final head
logits = self.head(features[0]) # only use the cls token
return logits
class MAE_Encoder_FeatureExtractor(torch.nn.Module):
'''
A feature extractor that extracts features from the encoder of the Masked Autoencoder.
'''
def __init__(self, encoder : MAE_Encoder) -> None:
super().__init__()
self.cls_token = encoder.cls_token
self.pos_embedding = encoder.pos_embedding
self.patchify = encoder.patchify
self.transformer = encoder.transformer
self.layer_norm = encoder.layer_norm
def forward(self, img):
patches = self.patchify(img)
patches = rearrange(patches, 'b c h w -> (h w) b c')
patches = patches + self.pos_embedding
patches = torch.cat([self.cls_token.expand(-1, patches.shape[1], -1), patches], dim=0)
patches = rearrange(patches, 't b c -> b t c')
features = self.layer_norm(self.transformer(patches))
# t is the number of patches, b is the batch size, c is the number of features
features = rearrange(features, 'b t c -> t b c')
return features
if __name__ == '__main__':
shuffle = PatchShuffle(0.75)
a = torch.rand(16, 2, 10)
b, forward_indexes, backward_indexes = shuffle(a)
print(b.shape)
img = torch.rand(2, 3, 32, 32)
encoder = MAE_Encoder()
decoder = MAE_Decoder()
features, backward_indexes = encoder(img)
print(forward_indexes.shape)
predicted_img, mask = decoder(features, backward_indexes)
print(predicted_img.shape)
loss = torch.mean((predicted_img - img) ** 2 * mask / 0.75)
print(loss)