Added the Visualizer, Transformer, MAE, and a test image

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+guineapig.jpg filter=lfs diff=lfs merge=lfs -text

ModelVisualizer.py

@@ -0,0 +1,89 @@
+import torch
+import numpy as np
+import torchvision.transforms as transforms
+
+import matplotlib.pyplot as plt
+from PIL import Image
+import maevit import MAEViT
+
+
+def visualize(model_path, img_path, figure_name):
+    
+    model = MAEViT(
+        image_size=224,
+        patch_size=16,
+        embed_dim=128,
+        encoder_layers=2,
+        encoder_heads=4,
+        mlp_ratio=2.0,
+        mask_ratio=0.75,
+        decoder_embed_dim=64,
+        decoder_layers=2,
+        decoder_heads=4,
+        dropout=0.1
+    )
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model.to(device)
+    checkpoint = torch.load(model_path, map_location=device)
+    model.load_state_dict(checkpoint)
+    model.eval()
+
+    
+    to_tensor = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std =[0.229, 0.224, 0.225]
+        )
+    ])
+
+    
+    img = Image.open(img_path).convert('RGB')
+    x = to_tensor(img).unsqueeze(0).to(device)  # [1,3,224,224]
+
+    
+    with torch.no_grad():
+        
+        x_enc, mask, ids_restore = model.forward_encoder(x)
+        
+        x_rec_patches = model.forward_decoder(x_enc, ids_restore)
+
+    
+    img_rec = model.unpatchify(x_rec_patches[:, 1:, :]) # exclude CLS    # [1,3,224,224]
+    img_patches = model.patchify(x)              # [1, num_patches, patch_dim]
+
+    masked_patches = img_patches.clone()
+    mask = mask.unsqueeze(-1).to(torch.bool)     # [1, num_patches, 1]
+    # masked_patches[mask] = 0
+    masked_patches = masked_patches.masked_fill(mask, 0)
+
+    img_masked = model.unpatchify(masked_patches) # [1,3,224,224]
+
+    inv_normalize = transforms.Normalize(
+        mean=[-m/s for m, s in zip((0.485,0.456,0.406),(0.229,0.224,0.225))],
+        std =[1/s for s in       (0.229,0.224,0.225)]
+    )
+    def to_img(tensor):
+        img = tensor.squeeze(0).cpu()
+        img = inv_normalize(img)
+        img = img.permute(1,2,0).clamp(0,1).numpy()
+        return img
+
+    orig_np     = to_img(x)
+    masked_np   = to_img(img_masked)
+    recon_np    = to_img(img_rec)
+
+    # 8. Plot
+    fig, axes = plt.subplots(1, 3, figsize=(15,5))
+    for ax, im, title in zip(axes,
+                             [orig_np, masked_np, recon_np],
+                             ['Original', 'Masked Input', 'Reconstruction']):
+        ax.imshow(im)
+        ax.set_title(title)
+        ax.axis('off')
+    plt.tight_layout()
+    plt.show()
+    plt.savefig(figure_name)
+
+visualize('MAE1.bin', img_path='guineapig.jpg', figure_name='figures/MAE_visualization1.png')

maevit.py

@@ -0,0 +1,246 @@
+from torch.utils.data import Dataset
+import torch.nn as nn
+from PIL import Image
+import json
+import os
+import random
+import torch
+import numpy as np
+from transformer import TransformerEncoder, TransformerEncoderLayer, TransformerDecoder, TransformerDecoderLayer
+    
+#includes both MAE and Vision Transformer for pretraining
+class MAEViT(nn.Module):
+    """
+    Masked Autoencoder (MAE) for Vision Transformer.
+    Encoder sees only a fraction of patches; decoder reconstructs all patches.
+    """
+    def __init__(
+        self,
+        # default values for ViT-B-16
+        image_size: int = 224,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        encoder_layers: int = 12,
+        encoder_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        mask_ratio: float = 0.75,
+        decoder_embed_dim: int = 512, 
+        decoder_layers: int = 8,
+        decoder_heads: int = 16,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        assert image_size % patch_size == 0, "Image size must be divisible by patch size"
+        self.in_chans = in_chans
+        self.image_size = image_size
+        self.patch_size = patch_size
+
+        #Conv2d trick to PATCHIFY AND EMBED (DIFFERENT FROM THE PATCHIFY Function
+        #which is used in validation)
+        self.conv_proj = nn.Conv2d(
+            in_channels = in_chans,
+            out_channels = embed_dim, #embed_dim is for the TOTAL; this is patch_dimen^2 * 3 (# of color channels)
+            kernel_size = patch_size, #this is so that the kernel is basically the patch (a square)
+            stride = patch_size #this ensures that the kernel moves so that the patches do not overlap
+        )
+        num_patches = (image_size // patch_size) ** 2 #just the number of patches since image_size // patch_size deals with only the dimension
+        self.mask_ratio = mask_ratio #75% is masked for best results with MAE
+
+        #set CLS token, a class token that contains a learnable vector that will eventually contain embeddings for the whole image
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) 
+        nn.init.normal_(self.cls_token, std = 0.02) #normal distribution
+
+        #Transformer encoder: learns contextual relationships b/t patches, generates embeddings
+        enc_layer = TransformerEncoderLayer(
+            embed_dim = embed_dim,
+            num_heads = encoder_heads, #for multihead attn
+            mlp_dim = int(embed_dim * mlp_ratio), 
+            dropout = dropout #used in MLP
+        )
+        self.encoder = TransformerEncoder(enc_layer, encoder_layers, embed_dim) #does self attn & feed forward
+
+        #Encoder -> Decoder (Linear Projection)
+        self.enc_to_dec = nn.Linear(embed_dim, decoder_embed_dim, bias = False)
+
+        #Decoder mask token (learnable placeholder token at each masked patch, helps decoder reconstruct those patches) and positional embedding generated
+        self.dec_mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim)) #mask tokens originally set to zero
+        self.dec_pos_embed = nn.Parameter(torch.empty(1, num_patches + 1, decoder_embed_dim)) #num_patches + 1 includes cls token
+        self.enc_pos_embed = nn.Parameter(torch.empty(1, num_patches + 1, embed_dim))
+        
+        nn.init.normal_(self.dec_mask_token, std=0.02)
+        nn.init.normal_(self.enc_pos_embed, std=0.02)
+        nn.init.normal_(self.dec_pos_embed, std=0.02)
+
+        #Decoder for the transformer: predicts the masked patches
+        dec_layer = TransformerDecoderLayer(
+            encoder_embed_dim = embed_dim,
+            decoder_embed_dim = decoder_embed_dim,
+            num_heads = decoder_heads,
+            mlp_dim = int(decoder_embed_dim * mlp_ratio),
+            dropout = dropout #a regularizer to prevent model from overfitting and possibly making decisions based on noise
+        )
+        self.decoder = TransformerDecoder(dec_layer, decoder_layers, embed_dim = decoder_embed_dim)
+
+        #Reconstruction for masked patches
+        self.pred = nn.Linear(decoder_embed_dim, (patch_size ** 2) * in_chans)
+
+        self.norm = nn.LayerNorm(embed_dim)
+
+    #patchify: converts image into tensors for patches
+    def patchify(self, imgs):
+        """
+        imgs: (B, C, H, W)
+        returns: (B, N, patch_size * patch_size * C)
+        """
+        B, C, H, W = imgs.shape
+        p = self.patch_size
+        assert H % p == 0 and W % p == 0, "Image dimensions must be divisible by the patch size."
+
+        h = H // p
+        w = W // p
+        patches = imgs.reshape(B, C, h, p, w, p)
+        patches = torch.einsum('nchawb->nhwabc', patches)
+        patches = patches.reshape(B, h * w, p * p * C)
+        #print("size of patches: ")
+        #print(patches.size())
+        return patches
+
+    #unpatchify: helps reconstruct image from patches (tensors -> images)
+    #is not actually needed, maybe for debugging
+    def unpatchify(self, x):
+        #x is a tensor of shape (B, num_patches, patch_size*patch_size*in_chans)
+        #x represents flattened pixel values
+        #imgs: (returned) has shape (B, in_chans, img_size, img_size)
+        patch_dimen = self.patch_size
+        h = int(x.shape[1]**0.5)
+        w = h
+        assert h * w == x.shape[1]
+        x = x.reshape(x.shape[0], h, w, patch_dimen, patch_dimen, self.in_chans)
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape([x.shape[0], self.in_chans, self.image_size, self.image_size])
+        return imgs
+
+    def random_masking(self, x):
+        """
+        Perform per-sample random masking by shuffling.
+        returns:
+            x_masked: Tensor with visible patches
+            mask: Tensor indicating which patches are visible (0) or masked (1)
+            ids_restore: Tensor to restore original order of patches
+        """
+        B, L, D = x.shape
+        #number of patches to keep
+        len_keep = int(L*(1 - self.mask_ratio))
+
+        #indices for visible patches by generating noise
+        noise = torch.rand(B, L, device=x.device)
+        ids_shuffle = torch.argsort(noise, dim=1)
+
+        #restore indices for unshuffling patches
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        #indices of kept patches
+        ids_keep = ids_shuffle[:, :len_keep]
+        #visible patches gathered
+        x_masked = torch.gather(x, dim=1, index = ids_keep.unsqueeze(-1).repeat(1, 1, D))
+        
+        #binary mask for patch embedding (1 is for masked, 0 is for visible)
+        mask = torch.ones(B, L, device=x.device)
+        #mask is unshuffled back into original patch order
+        mask[:, :len_keep] = 0 ## DONGHEE: THIS PART WAS MISSING IN THE ORIGINAL CODE
+        mask = torch.gather(mask, 1, ids_restore)
+
+        return x_masked, mask, ids_restore, ids_keep  
+    
+    def forward_encoder(self, imgs):
+
+        # 1. Patch embedding
+        x = self.conv_proj(imgs)            # [B, embed_dim, H/ps, W/ps]
+        x = x.flatten(2).transpose(1, 2)    # [B, N, embed_dim]
+        x = self.norm(x)                  # [B, N, embed_dim]
+        B, N, D = x.shape
+
+        # 2. Add positional embeddings (w/o class token)
+        #print(x.shape)
+        #print(self.enc_pos_embed.shape)
+        x = x + self.enc_pos_embed[:, 1:, :] 
+
+        # 3. Random masking
+        x_masked, mask, ids_restore, ids_keep = self.random_masking(x)
+
+        # 4. Encoder input (cls token + visible patches)
+        cls_token = self.cls_token + self.enc_pos_embed[:, :1, :]  # class token with positional embedding
+        cls_tokens = cls_token.expand(B, -1, -1)  # repeat for batch size
+        x_enc = torch.cat([cls_tokens, x_masked], dim=1)
+
+        # 5. Encoder forward
+        x_enc = self.encoder(x_enc) # TO DO
+
+        return x_enc, mask, ids_restore
+    
+    def forward_decoder(self, x_enc, ids_restore):
+        # encoder output needs to be projected to decoder embedding space
+        x_dec = self.enc_to_dec(x_enc)
+
+        #sequence unshuffled to original order
+        B, L, D = x_dec.shape
+        mask_tokens = self.dec_mask_token.repeat(B, ids_restore.shape[1] + 1 - x_dec.shape[1], 1)
+        
+        #concatenate output from heads?
+        x_no_cls = torch.cat([x_dec[:, 1:, :], mask_tokens], dim=1)
+        x_no_cls = torch.gather(x_no_cls, 1, ids_restore.unsqueeze(-1).repeat(1, 1, D))
+        x_dec = torch.cat([x_dec[:, :1, :], x_no_cls], dim=1)
+
+        #add positional embeddings
+        x_dec = x_dec + self.dec_pos_embed[:, :x_dec.size(1), :]
+
+        #decoder forward
+        x_dec = self.decoder(x_dec, x_enc) 
+
+        #predict pixels (without class token)
+        x_rec = self.pred(x_dec)
+
+        return x_rec
+    
+    def compute_mae_loss(self, imgs, pred, mask):
+        """
+        Mean Squared Error loss for masked patches
+        imgs: [N, 3, H, W]
+        pred: [N, L, p*p*3]
+        mask: [N, L], 0 is keep, 1 is remove, 
+        """
+
+        #mask: binary mask tensor
+        target = self.patchify(imgs) 
+        #print("target size: ")
+        #print(target.size())
+        #print("pred size: ")
+        #print(pred.size())
+
+        pred = pred[:, 1:, :]
+        loss = (pred - target)**2
+        loss = loss.mean(dim=-1)
+        #we don't want to calculate loss on visible patches, only masked patches
+        loss = (loss * mask).sum() / (mask.sum() + 1e-6)
+
+        return loss
+    
+    def forward(self, imgs):
+        """
+        Forward pass for MAE: encode, decode, and compute reconstruction loss.
+        imgs: [B, 3, H, W]
+        returns: reconstruction loss
+        """
+
+        # 1. Forward encoder
+        x_enc, mask, ids_restore = self.forward_encoder(imgs)
+
+        #x_enc = self.enc_to_dec(x_enc)
+
+        # 2. Forward decoder
+        x_rec = self.forward_decoder(x_enc, ids_restore)
+
+        loss = self.compute_mae_loss(imgs, x_rec, mask)
+
+        return loss

transformer.py

@@ -0,0 +1,239 @@
+import torch
+import torch.nn as nn
+import copy
+import math
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention:
+    1) Linear projections for Q, K, V
+    2) Scaled dot-product attention per head
+    3) Concatenate heads and final linear projection
+    https://arxiv.org/pdf/1706.03762
+    """
+    def __init__(self, embed_dim:int, key_dim:int, num_heads: int, dropout: float = 0.0):
+        super().__init__()
+        
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.embed_dim = embed_dim
+        self.key_dim = key_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scale = math.sqrt(self.head_dim) # square root of dk for scaling
+
+        # Separate projections for query, key, and value 3 diff transformations
+        # HINT: Linear projections for Q, K, V 
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(key_dim, embed_dim) #To Do
+        self.v_proj = nn.Linear(key_dim, embed_dim) #To Do
+        
+        # Output projection after concatenating heads (embed_dim -> embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim) #To Do
+        
+        # Dropouts
+        self.attn_dropout = nn.Dropout(dropout)
+        self.proj_dropout = nn.Dropout(dropout)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            query: [batch, seq_q, embed_dim]
+            key:   [batch, seq_k, embed_dim]
+            value: [batch, seq_k, embed_dim]
+        Returns:
+            out:   [batch, seq_q, embed_dim]
+        """
+        B, seq_q, _ = query.size()
+        _, seq_k, _ = key.size()
+
+        # 1) Project inputs and split into heads
+        q = self.q_proj(query).view(B, seq_q, self.num_heads, self.head_dim).transpose(1, 2)  # [B, heads, seq_q, head_dim]
+        k = self.k_proj(key).view(B, seq_k, self.num_heads, self.head_dim).transpose(1,2)    # [B, heads, seq_k, head_dim]
+        v = self.v_proj(value).view(B, seq_k, self.num_heads, self.head_dim).transpose(1,2)  # [B, heads, seq_k, head_dim]
+
+        # 2) Compute scaled dot-product attention
+        scores = (q @ k.transpose(-1, -2)) / self.scale # TO DO multiply q and k, then scale # [B, heads, seq_q, seq_k] swaps last twp dims?
+        weights = torch.softmax(scores, dim=-1) # TO DO apply softmax to scores
+        weights = self.attn_dropout(weights)
+        attn = weights @ v # TO DO multiply weights and v # [B, heads, seq_q, head_dim]
+
+        # 3) Concatenate heads
+        attn = attn.transpose(1, 2).contiguous().view(B, seq_q, self.embed_dim)  # [B, seq_q, embed_dim]
+
+        # 4) Final projection
+        out =  self.out_proj(attn) # TO DO apply output projection
+        out = self.proj_dropout(out)
+
+        return out
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    Transformer Encoder Layer:
+    1) Multi-head self-attention
+    2) Feed-forward network
+    3) Residual connections + LayerNorm
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        mlp_dim: int,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        
+        # 1) Self-attention
+        self.self_attn = MultiHeadAttention(
+            embed_dim=embed_dim,
+            key_dim=embed_dim,  # self-attention uses same dimension for Q, K, V
+            num_heads=num_heads,
+            dropout=dropout,
+        )
+
+        # 2) Feed-forward network using nn.Sequential
+        self.ffn = nn.Sequential(
+            nn.Linear(embed_dim, mlp_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_dim, embed_dim),
+        )
+
+        # 3) LayerNorm and Dropouts for residuals
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        # self.norm3 = nn.LayerNorm(embed_dim)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.ff_dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch, seq_len, embed_dim]
+        Returns:
+            Tensor of same shape
+        """
+        # 1) Self-attention block
+        x_norm = self.norm1(x)  # Normalize input
+        attn_out =  self.self_attn(x_norm, x_norm, x_norm)
+        x = x + self.attn_dropout(attn_out) # Residual connection
+
+        # 2) Feed-forward block
+        x_norm_ff = self.norm2(x)
+        ff = self.ffn(x_norm_ff)
+        x = x + self.ff_dropout(ff) # Residual connection
+        # x = self.norm3(x)
+
+        return x
+    
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, encoder_layer: TransformerEncoderLayer, num_layers: int, embed_dim: int):
+        super().__init__()
+
+        # Clone the provided encoder_layer num_layers times
+        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)])
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        
+        # TO DO pass through each layer
+        # HINT: Pass input through each encoder layer to update x
+        for layer in self.layers:
+            x = layer(x)
+
+        # Apply final normalization   
+        x = self.norm(x)
+        
+        return x
+
+
+class TransformerDecoderLayer(nn.Module):
+    """
+    Transformer Decoder Layer:
+    1) Self-attention
+    2) Cross-attention (over encoder features)
+    3) Feed-forward network
+    4) Residual connections + LayerNorm
+    """
+    def __init__(
+        self,
+        encoder_embed_dim: int,
+        decoder_embed_dim: int,
+        num_heads: int,
+        mlp_dim: int,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        
+        # 1. Self-attention on decoder input
+        self.self_attn = MultiHeadAttention(decoder_embed_dim, decoder_embed_dim, num_heads, dropout)
+
+        # 2. Cross-attention over encoder features
+        self.cross_attn = MultiHeadAttention(decoder_embed_dim, encoder_embed_dim, num_heads, dropout)
+
+        # 3. Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(decoder_embed_dim, mlp_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_dim, decoder_embed_dim)
+        )
+
+        # 4. LayerNorms and Dropouts
+        self.norm1 = nn.LayerNorm(decoder_embed_dim)
+        self.norm2 = nn.LayerNorm(decoder_embed_dim)
+        self.norm3 = nn.LayerNorm(decoder_embed_dim)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, encoder_features: torch.Tensor) -> torch.Tensor:
+        
+        # 1) Self-attention block
+        x_norm = self.norm1(x)  # Normalize input
+        sa = self.self_attn(x_norm, x_norm, x_norm)
+        # TO DO
+        x = x + self.dropout1(sa)
+        
+        # 2) Cross-attention block
+        x_norm_ca = self.norm2(x)
+        ca = self.cross_attn(x_norm_ca, encoder_features, encoder_features)
+        # TO DO
+        x = x + self.dropout2(ca)
+        
+        # 3) Feed-forward block
+        x_norm_ff = self.norm3(x)
+        ff = self.ffn(x_norm_ff)
+        # TO DO
+        x = x + self.dropout3(ff)
+
+        return x
+
+
+class TransformerDecoder(nn.Module):
+    
+    def __init__(
+        self,
+        decoder_layer: TransformerDecoderLayer,
+        num_layers: int,
+        embed_dim: int,
+    ):
+        super().__init__()
+
+        # Clone the provided decoder_layer num_layers times
+        self.layers = nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(num_layers)])
+
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor, encoder_features: torch.Tensor) -> torch.Tensor:
+        
+        # TODO pass through each layer
+        # HINT: Pass input through each encoder layer to update x
+        for layer in self.layers:
+            x = layer(x, encoder_features)
+        
+        x = self.norm(x)
+        
+        return x