Upload 3 files

load_celebA.py

@@ -0,0 +1,39 @@
+"""
+CelebFaces Attributes (CelebA) Dataset
+https://www.kaggle.com/datasets/jessicali9530/celeba-dataset
+"""
+
+
+import os
+
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+
+
+class CelebADataset(Dataset):
+
+    def __init__(self, root, img_shape=(64, 64)) -> None:
+        super().__init__()
+        self.root = root
+        self.img_shape = img_shape
+        self.filenames = sorted(os.listdir(root))
+
+    def __len__(self) -> int:
+        return len(self.filenames)
+
+    def __getitem__(self, index: int):
+        path = os.path.join(self.root, self.filenames[index])
+        img = Image.open(path).convert('RGB')
+        pipeline = transforms.Compose([
+            transforms.CenterCrop(168),
+            transforms.Resize(self.img_shape),
+            transforms.ToTensor()
+        ])
+        return pipeline(img)
+
+
+def get_dataloader(root='data/celebA/img_align_celeba', **kwargs):
+    dataset = CelebADataset(root, **kwargs)
+    return DataLoader(dataset, 16, shuffle=True)

main.py

@@ -0,0 +1,84 @@
+from time import time
+
+import torch
+import torch.nn.functional as F
+from torchvision.transforms import ToPILImage
+
+from dldemos.VAE.load_celebA import get_dataloader
+from dldemos.VAE.model import VAE
+
+# Hyperparameters
+n_epochs = 10
+kl_weight = 0.00025
+lr = 0.005
+
+
+def loss_fn(y, y_hat, mean, logvar):
+    recons_loss = F.mse_loss(y_hat, y)
+    kl_loss = torch.mean(
+        -0.5 * torch.sum(1 + logvar - mean**2 - torch.exp(logvar), 1), 0)
+    loss = recons_loss + kl_loss * kl_weight
+    return loss
+
+
+def train(device, dataloader, model):
+    optimizer = torch.optim.Adam(model.parameters(), lr)
+    dataset_len = len(dataloader.dataset)
+
+    begin_time = time()
+    # train
+    for i in range(n_epochs):
+        loss_sum = 0
+        for x in dataloader:
+            x = x.to(device)
+            y_hat, mean, logvar = model(x)
+            loss = loss_fn(x, y_hat, mean, logvar)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            loss_sum += loss
+        loss_sum /= dataset_len
+        training_time = time() - begin_time
+        minute = int(training_time // 60)
+        second = int(training_time % 60)
+        print(f'epoch {i}: loss {loss_sum} {minute}:{second}')
+        torch.save(model.state_dict(), 'dldemos/VAE/model.pth')
+
+
+def reconstruct(device, dataloader, model):
+    model.eval()
+    batch = next(iter(dataloader))
+    x = batch[0:1, ...].to(device)
+    output = model(x)[0]
+    output = output[0].detach().cpu()
+    input = batch[0].detach().cpu()
+    combined = torch.cat((output, input), 1)
+    img = ToPILImage()(combined)
+    img.save('work_dirs/tmp.jpg')
+
+
+def generate(device, model):
+    model.eval()
+    output = model.sample(device)
+    output = output[0].detach().cpu()
+    img = ToPILImage()(output)
+    img.save('work_dirs/tmp.jpg')
+
+
+def main():
+    device = 'cuda:0'
+    dataloader = get_dataloader()
+
+    model = VAE().to(device)
+
+    # If you obtain the ckpt, load it
+    model.load_state_dict(torch.load('dldemos/VAE/model.pth', 'cuda:0'))
+
+    # Choose the function
+    train(device, dataloader, model)
+    reconstruct(device, dataloader, model)
+    generate(device, model)
+
+
+if __name__ == '__main__':
+    main()

model.py

@@ -0,0 +1,91 @@
+"""
+Full definition of a VAE model, all of it in this single file.
+References:
+1) An Introduction to Variational Autoencoders:
+https://arxiv.org/abs/1906.02691
+"""
+
+
+import torch
+import torch.nn as nn
+
+
+class VAE(nn.Module):
+    """VAE for 64x64 face generation.
+
+    The hidden dimensions can be tuned.
+    """
+
+    def __init__(self, hiddens=[16, 32, 64, 128, 256], latent_dim=128) -> None:
+        super().__init__()
+
+        # encoder
+        prev_channels = 3
+        modules = []
+        img_length = 64
+        for cur_channels in hiddens:
+            modules.append(
+                nn.Sequential(
+                    nn.Conv2d(prev_channels,
+                              cur_channels,
+                              kernel_size=3,
+                              stride=2,
+                              padding=1), nn.BatchNorm2d(cur_channels),
+                    nn.ReLU()))
+            prev_channels = cur_channels
+            img_length //= 2
+        self.encoder = nn.Sequential(*modules)
+        self.mean_linear = nn.Linear(prev_channels * img_length * img_length,
+                                     latent_dim)
+        self.var_linear = nn.Linear(prev_channels * img_length * img_length,
+                                    latent_dim)
+        self.latent_dim = latent_dim
+
+        # decoder
+        modules = []
+        self.decoder_projection = nn.Linear(
+            latent_dim, prev_channels * img_length * img_length)
+        self.decoder_input_chw = (prev_channels, img_length, img_length)
+        for i in range(len(hiddens) - 1, 0, -1):
+            modules.append(
+                nn.Sequential(
+                    nn.ConvTranspose2d(hiddens[i],
+                                       hiddens[i - 1],
+                                       kernel_size=3,
+                                       stride=2,
+                                       padding=1,
+                                       output_padding=1),
+                    nn.BatchNorm2d(hiddens[i - 1]), nn.ReLU()))
+        modules.append(
+            nn.Sequential(
+                nn.ConvTranspose2d(hiddens[0],
+                                   hiddens[0],
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                nn.BatchNorm2d(hiddens[0]), nn.ReLU(),
+                nn.Conv2d(hiddens[0], 3, kernel_size=3, stride=1, padding=1),
+                nn.ReLU()))
+        self.decoder = nn.Sequential(*modules)
+
+    def forward(self, x):
+        encoded = self.encoder(x)
+        encoded = torch.flatten(encoded, 1)
+        mean = self.mean_linear(encoded)
+        logvar = self.var_linear(encoded)
+        eps = torch.randn_like(logvar)
+        std = torch.exp(logvar / 2)
+        z = eps * std + mean
+        x = self.decoder_projection(z)
+        x = torch.reshape(x, (-1, *self.decoder_input_chw))
+        decoded = self.decoder(x)
+
+        return decoded, mean, logvar
+
+    def sample(self, device='cuda'):
+        z = torch.randn(1, self.latent_dim).to(device)
+        x = self.decoder_projection(z)
+        x = torch.reshape(x, (-1, *self.decoder_input_chw))
+        decoded = self.decoder(x)
+        return decoded