Upload 7 files

RoPE.py

@@ -0,0 +1,22 @@
+import torch
+from einops import rearrange
+
+def generate_angles_2d(H,W,D, freq=None):
+    freq = torch.tensor([10000**(-2*i/D) for i in range(int(D/2))]) if freq is None else freq
+    pos = torch.outer(torch.linspace(-1, 1, steps=H),torch.linspace(-1, 1, steps=W))
+    freq_tensor = torch.einsum("ij,k->ijk", pos, freq)
+    return freq_tensor
+
+def apply_angles_2d(x, f):
+    x_reshaped = rearrange(x, "B h H W (D p) -> B h H W D p", p=2)
+    real = x_reshaped[..., 0]
+    imag = x_reshaped[..., 1]
+    cosines, sines = f.cos(), f.sin()
+    # r , i -> rcos-isin , rsin icos
+    rot_real = real * cosines - imag * sines
+    rot_imag = real * sines + imag * cosines
+    rot_full = torch.concat((rot_real.unsqueeze(-1), rot_imag.unsqueeze(-1)), dim=-1)
+    return rearrange(rot_full, "B h H W D p -> B h H W (D p)", p=2)
+
+# Sanity Check :)
+print(apply_angles_2d(torch.randn(1,8,64,64,768), generate_angles_2d(64,64,768)).shape)

attention.py

@@ -0,0 +1,63 @@
+from RoPE import apply_angles_2d, generate_angles_2d
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class Attention(nn.Module):
+    def __init__(self, H,W, emb_dim, n_heads=8):
+        super().__init__()
+        self.H = H
+        self.W = W
+        self.n_heads = n_heads
+        head_dim = emb_dim // n_heads
+        self.qkv = nn.Linear(emb_dim, 3*emb_dim, bias=False)
+        self.apply_angles_2d = apply_angles_2d
+        self.proj = nn.Linear(emb_dim, emb_dim)
+        self.register_buffer("freq", generate_angles_2d(H, W, head_dim), persistent=False)
+
+    def forward(self, x):
+        B, N, D = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+
+        # to 2D
+        q = rearrange(q, "B (H W) (h D) -> B h H W D", H=self.H, W=self.W, h=self.n_heads)
+        k = rearrange(k, "B (H W) (h D) -> B h H W D", H=self.H, W=self.W, h=self.n_heads)
+        v = rearrange(v, "B (H W) (h D) -> B h H W D", H=self.H, W=self.W, h=self.n_heads)
+
+        q = apply_angles_2d(q, self.freq)
+        k = apply_angles_2d(k, self.freq)
+        v = apply_angles_2d(v, self.freq)
+
+        # to 1D
+        q = rearrange(q, "B h H W D -> B h (H W) D", H=self.H, W=self.W, h=self.n_heads)
+        k = rearrange(k, "B h H W D -> B h (H W) D", H=self.H, W=self.W, h=self.n_heads)
+        v = rearrange(v, "B h H W D -> B h (H W) D", H=self.H, W=self.W, h=self.n_heads)
+
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "B h N D -> B N (h D)")
+        x = self.proj(x)
+        return x
+
+class ViTBlock(nn.Module):
+  def __init__(self, H, W, emb_dim, n_heads=8, dropout=0.1):
+    self.H, self.W, self.emb_dim = H, W, emb_dim
+    super().__init__()
+    self.attn = nn.Sequential(nn.LayerNorm(emb_dim),
+                              Attention(H,W,emb_dim,n_heads=n_heads))
+    self.MLP = nn.Sequential(nn.LayerNorm(emb_dim),
+                             nn.Linear(emb_dim, emb_dim*4, bias=True),
+                             nn.GELU(),
+                             nn.Dropout(dropout),
+                             nn.Linear(emb_dim*4, emb_dim, bias=True),
+                             nn.Dropout(dropout))
+  def forward(self, x):
+    assert x.ndim == 3, f"Expected shape [B, N, D], but got shape {x.shape}. You probably passed [B, H, W, D] instead."
+    assert x.shape == torch.Size([x.shape[0], self.H * self.W, self.emb_dim]), f"Expected shape [B, N, D] -> {torch.Size([x.shape[0], self.H * self.W, self.emb_dim])}, got {x.shape}"
+    x = x + self.attn(x)
+    x = x + self.MLP(x)
+    return x
+  
+# Sanity Check :)
+print(ViTBlock(64,64,384)(torch.randn(1, 64**2, 384)).shape)

autoencoder.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from attention import ViTBlock
+
+# Global Parameters
+image_shape = 256
+emb_dim = 768
+patch_size = 16
+
+class Encoder(nn.Module):
+  def __init__(self, latent_dim, image_shape=image_shape, emb_dim=emb_dim, patch_size=patch_size, n_heads=8, dropout=0.1, layers=6, gaussian=False):
+    super().__init__()
+    self.patchifier = nn.Conv2d(3, emb_dim, kernel_size=patch_size, stride=patch_size)
+    self.Blocks = nn.ModuleList([ViTBlock(image_shape // patch_size, image_shape // patch_size, emb_dim, n_heads=8, dropout=dropout) for _ in range(layers)])
+    self.ln = nn.LayerNorm(emb_dim)
+    self.compress_latent = nn.Linear(emb_dim, latent_dim)
+
+  def forward(self,x):
+    x = self.patchifier(x)
+    x = rearrange(x, "B D H W -> B (H W) D") # Flatten to B, N, D
+    for vitBlock in self.Blocks:
+      x = vitBlock(x)
+    x = self.ln(x)
+    x = self.compress_latent(x)
+    return x
+
+class Decoder(nn.Module):
+  def __init__(self, latent_dim, image_shape=image_shape, emb_dim=emb_dim, patch_size=patch_size, n_heads=8, dropout=0.1, layers=6, gaussian=False):
+    super().__init__()
+    self.hw = image_shape // patch_size
+    self.patch_size = patch_size
+    self.decompress_latent = nn.Linear(latent_dim, emb_dim)
+    self.ln = nn.LayerNorm(emb_dim)
+    self.emb_to_patch = nn.Linear(emb_dim, 3*(patch_size**2))
+    self.Blocks = nn.ModuleList([ViTBlock(image_shape // patch_size, image_shape // patch_size, emb_dim, n_heads=8, dropout=dropout) for _ in range(layers)])
+
+  def forward(self,x):
+    x = self.decompress_latent(x)
+    for vitBlock in self.Blocks:
+      x = vitBlock(x)
+    self.ln(x)
+    #shape is [B HW/p**2 (3 p p)]
+    x = self.emb_to_patch(x)
+    assert x.shape == torch.Size([x.shape[0], self.hw**2, 3*(self.patch_size**2)]), f"Expected shape {torch.Size([x.shape[0], self.hw**2, 3*(self.patch_size**2)])} got {x.shape}"
+    x = rearrange(x, "B (H W) (D p1 p2) -> B D (H p1) (W p2)", H=self.hw, W=self.hw, p1=self.patch_size, p2=self.patch_size) # Expand to B, H, W, D
+    return F.tanh(x)

autoencoder_test.py

@@ -0,0 +1,31 @@
+from autoencoder import Encoder, Decoder
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import cv2
+
+image_shape = 256
+emb_dim = 768
+patch_size = 16
+
+encoder = Encoder(latent_dim=16, 
+                  image_shape=image_shape, 
+                  emb_dim=emb_dim, 
+                  patch_size=patch_size)
+encoder.load_state_dict(torch.load("encoder16.pt", map_location=torch.device('cpu')))
+
+decoder = Decoder(latent_dim=16,
+                  image_shape=image_shape, 
+                  emb_dim=emb_dim, 
+                  patch_size=patch_size)
+decoder.load_state_dict(torch.load("decoder16.pt", map_location=torch.device('cpu')))
+
+image = cv2.imread("test_image.jpg")
+image = cv2.resize(image, (image_shape, image_shape))
+image = torch.tensor(image, dtype=torch.float32, device='cpu').permute(2, 0, 1) / 127.5 - 1.0
+image = image.unsqueeze(0) 
+with torch.no_grad():
+    z = encoder(image)
+    x = decoder(z)
+plt.imshow(x[0].permute(1, 2, 0).numpy()*0.5 + 0.5)
+plt.show()

objectives.py

@@ -0,0 +1,55 @@
+import torch.nn as nn
+from torchvision.models import vgg16, VGG16_Weights
+
+class Discriminator(nn.Module):
+    def __init__(self, img_shape, filters=[256,512]):
+        super().__init__()
+        module_list = [nn.Conv2d(img_shape[0], filters[0], kernel_size=3, stride=2, padding=1),
+                       nn.BatchNorm2d(filters[0]),
+                       nn.LeakyReLU(0.2)]
+        for i in range(1,len(filters)):
+            module_list += [nn.Conv2d(filters[i-1], filters[i], kernel_size=3, stride=2, padding=1),
+                            nn.BatchNorm2d(filters[i]),
+                            nn.LeakyReLU(0.2)]
+
+        self.convs = nn.Sequential(*module_list)
+        self.mlp = nn.Sequential(nn.Conv2d(filters[-1], 1, kernel_size=1, stride=1, padding=0))
+
+    def forward(self, x):
+        x = self.convs(x)
+        x = self.mlp(x)
+        return x
+
+class vgg_builder(nn.Module):
+    def __init__(self):
+        super(vgg_builder, self).__init__()
+        convs = vgg16(weights=VGG16_Weights.IMAGENET1K_V1).features
+        self.N_slices = 5
+        self.slices = nn.ModuleList(list(nn.Sequential() for _ in range(self.N_slices)))
+        for x in range(4):
+            self.slices[0].add_module(str(x), convs[x])
+        for x in range(4, 9):
+            self.slices[1].add_module(str(x), convs[x])
+        for x in range(9, 16):
+            self.slices[2].add_module(str(x), convs[x])
+        for x in range(16, 23):
+            self.slices[3].add_module(str(x), convs[x])
+        for x in range(23, 30):
+            self.slices[4].add_module(str(x), convs[x])
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, x):
+        feat_map = []
+        x = (x+1)/2
+        x = self.slices[0](x)
+        feat_map.append(x)
+        x = self.slices[1](x)
+        feat_map.append(x)
+        x = self.slices[2](x)
+        feat_map.append(x)
+        x = self.slices[3](x)
+        feat_map.append(x)
+        x = self.slices[4](x)
+        feat_map.append(x)
+        return feat_map

train.py

@@ -0,0 +1,81 @@
+import kagglehub
+import cv2
+import os
+from IPython.display import clear_output
+import torch
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+import torch.nn as nn
+import matplotlib.pyplot as plt
+from autoencoder import Encoder, Decoder
+from trainer import Trainer
+from objectives import Discriminator, vgg_builder
+
+# Global Parameters
+image_shape = 256
+emb_dim = 768
+patch_size = 16
+
+image_path = kagglehub.dataset_download("awsaf49/coco-2017-dataset")
+data = []
+for dirpath, _, filenames in os.walk(image_path):
+    for filename in filenames:
+        if filename.endswith("jpg"):
+            name = os.path.join(dirpath, filename)
+            img = cv2.imread(name)
+            img = cv2.resize(img, (image_shape,image_shape))
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = img.astype(np.float32) / 127.5 - 1.0
+            img = torch.tensor(img).permute(2,0,1)
+            data.append(img)
+            clear_output(wait=1)
+            print(f"{len(data)/1670:.2f}%")
+print(len(data))
+
+class CustomDataset(Dataset):
+    def __init__(self, data):
+        self.indices = np.arange(len(data))
+        np.random.shuffle(self.indices)
+        self.data = data
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        return torch.tensor(self.data[self.indices[idx]], dtype=torch.float32)
+
+# Sanity Check :)
+plt.imshow(CustomDataset(data)[0].permute(1,2,0)/2+0.5)
+
+encoder = Encoder(latent_dim=16)
+decoder = Decoder(latent_dim=16)
+D = Discriminator((3,256,256))
+
+vgg = vgg_builder()
+for param in vgg.parameters():
+  param.requires_grad = False
+vgg.eval()
+
+print(f"encoder: {sum(p.numel() for p in encoder.parameters())/(262144):.3f}MB")
+print(f"decoder: {sum(p.numel() for p in decoder.parameters())/(262144):.3f}MB")
+print(f"Discriminator: {sum(p.numel() for p in D.parameters())/(262144):.3f}MB")
+print(f"VGG: {sum(p.numel() for p in vgg.parameters())/(262144):.3f}MB")
+
+batch_size = 16
+dataset    = CustomDataset(data)
+loader     = DataLoader(dataset,
+                        batch_size=batch_size,
+                        shuffle=True,
+                        num_workers=8,
+                        pin_memory=True)
+epochs = 5
+trainer = Trainer(encoder, decoder, D, vgg, ["mse", "gan", "vgg", "KL"], len(loader) if "loader" in locals() else 0, isViT=1)
+for epoch in range(1, epochs):
+    index = 0
+    for i, x in enumerate(loader):
+        trainer.train_step(x, freeze_disc=0, with_mse=1, freeze_ae=0)
+    trainer.update_epoch()
+
+torch.save(encoder.state_dict(), "encoder16.pt")
+torch.save(decoder.state_dict(), "decoder16.pt")
+torch.save(D.state_dict(), "discriminator16.pt")

trainer.py

@@ -0,0 +1,98 @@
+import torch
+import torch.nn.functional as F
+from IPython.display import clear_output
+
+# @title Trainer
+class Trainer():
+    def __init__(self, encoder, decoder, D, vgg, losses, data_len, ema=3, a_disc=1, a_vae=1, a_KL=0.1, isViT=True):
+        self.vgg_schedule = None
+        self.ema = 2/(ema+1)
+        self.a_disc = a_disc
+        self.a_vae = a_vae
+        self.a_KL = a_KL
+
+        self.isViT = isViT
+        self.encoder = encoder
+        self.decoder = decoder
+        self.D = D
+        self.vgg = vgg
+        self.encoder_optimizer = torch.optim.Adam(self.encoder.parameters(),  lr=1e-5)
+        self.encoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.encoder_optimizer, T_max=50)
+        self.decoder_optimizer = torch.optim.Adam(self.decoder.parameters(),  lr=1e-5)
+        self.decoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.decoder_optimizer, T_max=50)
+        self.D_optimizer = torch.optim.Adam(self.D.parameters(),  lr=4e-5)
+        self.D_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.D_optimizer, T_max=50)
+        self.losses = losses
+        self.loss_vals = {loss:0 for loss in losses}
+        self.data_len = data_len
+        self.loss_record = []
+        self.epoch = 1
+        self.index = 1
+        self.device = torch.device("cuda")
+
+        self.encoder.to(self.device)
+        self.decoder.to(self.device)
+        self.D.to(self.device)
+        self.vgg.to(self.device)
+
+    def train_step(self, x, with_mse=False, freeze_ae=False, freeze_disc=False):
+        self.index += 1
+        x = x.to(self.device)
+        with torch.no_grad():
+            x_hat = self.decoder(self.encoder(x.permute(0,2,3,1))).permute(0,3,1,2) if not self.isViT else self.decoder(self.encoder(x))
+        if not freeze_disc:
+            disc_loss = F.relu(1. - self.D(x)).mean() + F.relu(1. + self.D(x_hat)).mean() # Hinge
+            self.D_optimizer.zero_grad()
+            disc_loss.backward()
+            self.D_optimizer.step()
+            self.D_scheduler.step()
+
+        if not freeze_ae:
+            z = self.encoder(x.permute(0,2,3,1)) if not self.isViT else self.encoder(x)
+            x_hat = self.decoder(z).permute(0,3,1,2) if not self.isViT else self.decoder(z)
+            mse = F.mse_loss(x_hat, x)
+            KL = 0.5 * (z.mean() ** 2)
+            vgg_real = self.vgg(x)
+            vgg_fake = self.vgg(x_hat)
+            vgg_loss = 0
+            for i in range(len(vgg_real)):
+                vgg_loss += F.mse_loss(vgg_real[i], vgg_fake[i]) 
+
+            adv_loss = 0
+            if not freeze_disc:
+                adv_loss = -(self.D(self.decoder(self.encoder(x))).mean())
+
+            loss = mse * with_mse + self.a_KL* KL + vgg_loss + self.a_vae * adv_loss
+            self.encoder_optimizer.zero_grad()
+            self.decoder_optimizer.zero_grad()
+            loss.backward()
+            self.encoder_optimizer.step()
+            self.decoder_optimizer.step()
+            self.encoder_scheduler.step()
+            self.decoder_scheduler.step()
+
+        self.update_batch({"mse":mse.item() if not freeze_ae else 0,
+                           "gan":disc_loss.item() if not freeze_disc else 0,
+                           "vgg":vgg_loss.item() if not freeze_ae else 0,
+                           "KL":z.mean() if not freeze_ae else 0})
+
+    def update_batch(self, loss_vals):
+        clear_output(wait=True)
+        for record in self.loss_record:
+            print(record)
+        self.loss_vals = {loss:(1-self.ema)*self.loss_vals[loss] + self.ema*loss_vals[loss] for loss in self.losses}
+        print(f"epoch:{self.epoch} ", end="")
+        for loss in self.losses:
+            print(f"{loss}: {self.loss_vals[loss]:.3f} ", end="")
+        for _ in range(int(self.index * 20 / self.data_len)):
+            print("=", end="")
+        for _ in range(int(self.index * 20 / self.data_len),20):
+            print("-", end="")
+
+    def update_epoch(self):
+        self.index = 0
+        record = f"epoch:{self.epoch} "
+        for loss in self.losses:
+            record += f"{loss}: {self.loss_vals[loss]:.3f} "
+        self.loss_record.append(record)
+        self.epoch += 1