First commit for VAE space

.gitignore

@@ -0,0 +1,11 @@
+medium/
+data/
+.vscode/
+__pycache__/
+.ipynb_checkpoints/
+lightning_logs/
+log_images/
+logs/
+*.env
+.idea/
+saved_models/

Dockerfile

@@ -0,0 +1,9 @@
+FROM python:3.8-slim-buster
+WORKDIR /app
+EXPOSE $PORT
+
+COPY requirements.txt /
+RUN pip3 install -r /requirements.txt
+COPY . /app
+
+CMD streamlit run app.py --server.port $PORT

README.md

@@ -1,12 +1,20 @@
----
-title: VAE
-emoji: 📊
-colorFrom: pink
-colorTo: red
-sdk: streamlit
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+### VAE with Pytorch-Lightning
+
+This is inspired from vae-playground. This is an example where we test out vae and conv_vae models with multiple datasets 
+like MNIST, celeb-a and MNIST-Fashion datasets.
+
+This also comes with an example streamlit app & deployed at huggingface.
+
+
+## Model Training
+
+You can train the VAE models by using `train.py` and editing the `config.yaml` file. \
+Hyperparameters to change are:
+- model_type [vae|conv_vae]
+- alpha
+- hidden_dim
+- dataset [celeba|mnist|fashion-mnist]
+
+There are other configurations that can be changed if required like height, width, channels etc. It also contains the pytorch-lightning configs as well.
+
+

app.py

@@ -0,0 +1,55 @@
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+import os
+import utils
+from PIL import Image
+
+
+st.set_page_config("VAE MNIST Pytorch Lightning")
+st.title("VAE Playground")
+# title_img = Image.open("images/title_img.jpg")
+
+# st.image(title_img)
+st.markdown(
+    "This is a simple streamlit app to showcase how a simple VAEs."
+)
+
+def load_model_files():
+    files = os.listdir("./saved_models/")
+    # Docker creates some whiteout files which mig
+    files = [i for i in files if ".ckpt" in i]
+    clean_names = [utils.parse_model_file_name(name) for name in files]
+    return {k: v for k, v in zip(clean_names, files)}
+
+
+file_name_map = load_model_files()
+files = list(file_name_map.keys())
+
+st.header("🖼️ Image Reconstruction", "recon")
+
+with st.form("reconstruction"):
+    model_name = st.selectbox("Choose Model:", files,
+                              key="recon_model_select")
+    recon_model_name = file_name_map[model_name]
+    recon_canvas = st_canvas(
+        # Fixed fill color with some opacity
+        fill_color="rgba(255, 165, 0, 0.3)",
+        stroke_width=8,
+        stroke_color="#FFFFFF",
+        background_color="#000000",
+        update_streamlit=True,
+        height=150,
+        width=150,
+        drawing_mode="freedraw",
+        key="recon_canvas",
+    )
+    submit = st.form_submit_button("Perform Reconstruction")
+    if submit:
+        recon_model = utils.load_model(recon_model_name)
+        inp_tens = utils.canvas_to_tensor(recon_canvas)
+        _, _, out = recon_model(inp_tens)
+        out = (out+1)/2
+        out_img = utils.resize_img(utils.tensor_to_img(out), 150, 150)
+if submit:
+    st.image(out_img)
+    

config.yaml

@@ -0,0 +1,20 @@
+training_params:
+  max_epochs: 30
+  auto_lr_find: false
+  gpus: 1
+model_params:
+  model_type: conv-vae # vae or conv-vae
+  lr: 0.005
+  batch_size: 1
+  hidden_size: 4096
+  latent_size: 128
+  alpha: 1024
+  dataset: "fashion-mnist"
+  save_images: true
+  save_path: "log_images/"
+  channels: 1
+  height: 64
+  width: 64
+logger_params:
+  name: "conv-vae"
+  save_dir: "logs/"

config/__init__.py

@@ -0,0 +1,2 @@
+from .config import config
+__all__ = ["config"]

config/config.py

@@ -0,0 +1,59 @@
+from pydantic import BaseModel
+from typing import Optional, Union
+import yaml
+
+
+class TrainConfig(BaseModel):
+    max_epochs: int
+    auto_lr_find: Union[bool, int]
+    gpus: int
+
+
+class VAEConfig(BaseModel):
+    model_type: str
+    hidden_size: int
+    latent_size: int
+    alpha: int
+    dataset: str
+    batch_size: Optional[int] = 64
+    save_images: Optional[bool] = False
+    lr: Optional[float] = None
+    save_path: Optional[str] = None
+
+
+class ConvVAEConfig(VAEConfig):
+    channels: int
+    height: int
+    width: int
+
+
+class LoggerConfig(BaseModel):
+    name: str
+    save_dir: str
+
+
+class Config(BaseModel):
+    model_config: Union[VAEConfig, ConvVAEConfig]
+    train_config: TrainConfig
+    model_type: str
+    log_config: LoggerConfig
+
+
+def load_config(path="config.yaml"):
+    config = yaml.load(open(path), yaml.SafeLoader)
+    model_type = config['model_params']['model_type']
+    if model_type == "vae":
+        model_config = VAEConfig(**config["model_params"])
+    elif model_type == "conv-vae":
+        model_config = ConvVAEConfig(**config["model_params"])
+    else:
+        raise NotImplementedError(f"Model {model_type} is not implemented")
+    train_config = TrainConfig(**config["training_params"])
+    log_config = LoggerConfig(**config["logger_params"])
+    config = Config(model_config=model_config, train_config=train_config,
+                    model_type=model_type, log_config=log_config)
+
+    return config
+
+
+config = load_config()

inference.py

@@ -0,0 +1,24 @@
+from models import vae_models
+from config import config
+from PIL import Image 
+from torchvision.transforms import Resize, ToPILImage, Compose
+
+from utils import load_model, tensor_to_img, resize_img, export_to_onnx
+
+
+
+def predict(model_ckpt="vae_alpha_1024_dim_128.ckpt"):
+    model_type = config.model_type
+    model = vae_models[model_type].load_from_checkpoint(f"./saved_models/{model_ckpt}")
+    model.eval()
+    test_iter = iter(model.test_dataloader())
+    d, _ = next(test_iter)
+    _, _, out = model(d)
+    out_img = tensor_to_img(out)
+    return out_img
+
+    
+
+if __name__ == "__main__":
+    predict()
+    # export_to_onnx("./saved_models/vae.ckpt")

models/__init__.py

@@ -0,0 +1,11 @@
+from .vae import VAE, Flatten, Stack  # noqa: F401
+from .conv_vae import Conv_VAE  # noqa: F401
+
+__all__ = [
+    'VAE', 'Flatten', 'Stack'
+    'Conv_VAE',
+]
+vae_models = {
+    "conv-vae": Conv_VAE,
+    "vae": VAE
+}

models/conv_vae.py

@@ -0,0 +1,239 @@
+from .vae import VAE, Flatten, Stack
+import torch.nn as nn
+import pytorch_lightning as pl
+import torch
+import os
+import random
+from typing import Optional
+import torchvision.transforms as transforms
+from torchvision.datasets import MNIST, FashionMNIST, CelebA
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+from torchvision.utils import save_image
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+
+class PrintShape(nn.Module):
+    def __init__(self):
+        super(PrintShape, self).__init__()
+
+    def forward(self, x):
+        # Do your print / debug stuff here
+        # print(f"Shape: {x.shape}")
+        return x
+
+class UnFlatten(nn.Module):
+    def forward(self, input, size=4096):
+        # print("Unflatteing")
+        return input.view(input.size(0), size, 1, 1)
+
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        # print("Flattening")
+        return input.view(input.size(0), -1)
+
+class Conv_VAE(pl.LightningModule):
+    def __init__(self, channels: int, height: int, width: int, lr: int,
+                 latent_size: int, hidden_size: int, alpha: int, batch_size: int,
+                 dataset: Optional[str] = None,
+                 save_images: Optional[bool] = None,
+                 save_path: Optional[str] = None, **kwargs):
+        super().__init__()
+        self.latent_size = latent_size
+        self.hidden_size = hidden_size
+        if save_images:
+            self.save_path = f'{save_path}/{kwargs["model_type"]}_images/'
+        self.save_hyperparameters()
+        self.save_images = save_images
+        self.lr = lr
+        self.batch_size = batch_size
+        self.alpha = alpha
+        self.dataset = dataset
+        assert not height % 4 and not width % 4, "Choose height and width to "\
+            "be divisible by 4"
+        self.channels = channels
+        self.height = height
+        self.width = width
+        self.latent_size = latent_size
+        self.save_hyperparameters()
+
+        self.data_transform = transforms.Compose([
+            transforms.Resize(64),
+            transforms.CenterCrop((64, 64)),
+            transforms.ToTensor()
+        ])
+
+
+        self.encoder = nn.Sequential(
+            PrintShape(),
+            nn.Conv2d(self.channels, 32, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(32),
+            nn.LeakyReLU(),
+            PrintShape(),
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(),
+            PrintShape(),
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(),
+            PrintShape(),
+            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(),
+            PrintShape(),
+            Flatten(),
+            PrintShape(),
+        )
+
+        self.fc1 = nn.Linear(self.hidden_size, self.latent_size)
+        self.fc2 = nn.Linear(self.latent_size, self.hidden_size)
+
+        self.decoder = nn.Sequential(
+            PrintShape(),
+            # nn.Linear(self.hidden_size, self.hidden_size),
+            # PrintShape(),
+            # nn.BatchNorm1d(self.hidden_size),
+            UnFlatten(),
+            PrintShape(),
+            nn.ConvTranspose2d(self.hidden_size, 256, kernel_size=6, stride=2, padding=1),
+            PrintShape(),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            PrintShape(),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            PrintShape(),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(32),
+            PrintShape(),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(32, self.channels, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(self.channels),
+            PrintShape(),
+            nn.Sigmoid(),
+        )
+
+    def encode(self, x):
+        hidden = self.encoder(x)
+        mu, log_var = self.fc1(hidden), self.fc1(hidden)
+        # print("Encoded")
+        return mu, log_var
+
+    def decode(self, z):
+        # print("Decoding")
+        # f = nn.Linear(self.latent_size, self.hidden_size)
+        z = self.fc2(z)
+        # print(f"L: {z.shape}")
+        x = self.decoder(z)
+        return x
+    
+    def reparametrize(self, mu, log_var):
+        # Reparametrization Trick to allow gradients to backpropagate from the
+        # stochastic part of the model
+        sigma = torch.exp(0.5*log_var)
+        z = torch.randn_like(sigma)
+        return mu + sigma*z
+
+    def training_step(self, batch, batch_idx):
+        x, _ = batch
+        mu, log_var, x_out = self.forward(x)
+        kl_loss = (-0.5*(1+log_var - mu**2 -
+                         torch.exp(log_var)).sum(dim=1)).mean(dim=0)
+        recon_loss_criterion = nn.MSELoss()
+        recon_loss = recon_loss_criterion(x, x_out)
+        # print(kl_loss.item(),recon_loss.item())
+        loss = recon_loss*self.alpha + kl_loss
+
+        self.log('train_loss', loss, on_step=False,
+                 on_epoch=True, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, _ = batch
+        mu, log_var, x_out = self.forward(x)
+
+        kl_loss = (-0.5*(1+log_var - mu**2 -
+                         torch.exp(log_var)).sum(dim=1)).mean(dim=0)
+        recon_loss_criterion = nn.MSELoss()
+        recon_loss = recon_loss_criterion(x, x_out)
+        # print(kl_loss.item(),recon_loss.item())
+        loss = recon_loss*self.alpha + kl_loss
+        self.log('val_kl_loss', kl_loss, on_step=False, on_epoch=True)
+        self.log('val_recon_loss', recon_loss, on_step=False, on_epoch=True)
+        self.log('val_loss', loss, on_step=False, on_epoch=True)
+        # print(x.mean(),x_out.mean())
+        return x_out, loss
+
+    def validation_epoch_end(self, outputs):
+        if not self.save_images:
+            return
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+        choice = random.choice(outputs)
+        output_sample = choice[0]
+        output_sample = output_sample.reshape(-1, 1, self.width, self.height)
+        # output_sample = self.scale_image(output_sample)
+        save_image(
+            output_sample,
+            f"{self.save_path}/epoch_{self.current_epoch+1}.png",
+            # value_range=(-1, 1)
+        )
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=(self.lr or self.learning_rate))
+        lr_scheduler = ReduceLROnPlateau(optimizer,)
+        return {
+            "optimizer": optimizer, "lr_scheduler": lr_scheduler,
+            "monitor": "val_loss"
+        }
+    
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        hidden = self.reparametrize(mu, log_var)
+        output = self.decode(hidden)
+        return mu, log_var, output
+
+    # Functions for dataloading
+    def train_dataloader(self):
+        if self.dataset == "mnist":
+            train_set = MNIST('data/', download=True,
+                              train=True, transform=self.data_transform)
+        elif self.dataset == "fashion-mnist":
+            train_set = FashionMNIST(
+                'data/', download=True, train=True,
+                transform=self.data_transform)
+        elif self.dataset == "celeba":
+            train_set = CelebA('data/', download=False, split="train", transform=self.data_transform)
+        return DataLoader(train_set, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        if self.dataset == "mnist":
+            val_set = MNIST('data/', download=True, train=False,
+                            transform=self.data_transform)
+        elif self.dataset == "fashion-mnist":
+            val_set = FashionMNIST(
+                'data/', download=True, train=False,
+                transform=self.data_transform)
+        elif self.dataset == "celeba":
+            val_set = CelebA('data/', download=False, split="valid", transform=self.data_transform)
+        return DataLoader(val_set, batch_size=self.batch_size)
+    
+    def test_dataloader(self):
+        if self.dataset == "mnist":
+            val_set = MNIST('data/', download=True, train=False,
+                            transform=self.data_transform)
+        elif self.dataset == "fashion-mnist":
+            val_set = FashionMNIST(
+                'data/', download=True, train=False,
+                transform=self.data_transform)
+        elif self.dataset == "celeba":
+            val_set = CelebA('data/', download=False, split="test", transform=self.data_transform)
+        return DataLoader(val_set, batch_size=self.batch_size)
+
+
+

models/vae.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+import random
+from torchvision.datasets import MNIST, FashionMNIST, CelebA
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+from torchvision.utils import save_image
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+
+import os
+from typing import Optional
+
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class Stack(nn.Module):
+    def __init__(self, channels, height, width):
+        super(Stack, self).__init__()
+        self.channels = channels
+        self.height = height
+        self.width = width
+
+    def forward(self, x):
+        return x.view(x.size(0), self.channels, self.height, self.width)
+
+
+class VAE(pl.LightningModule):
+    def __init__(self, latent_size: int, hidden_size: int, alpha: int, lr: float,
+                 batch_size: int,
+                 dataset: Optional[str] = None,
+                 save_images: Optional[bool] = None,
+                 save_path: Optional[str] = None, **kwargs):
+        """Init function for the VAE
+
+        Args:
+
+        latent_size (int): Latent Hidden Size
+        alpha (int): Hyperparameter to control the importance of
+        reconstruction loss vs KL-Divergence Loss
+        lr (float): Learning Rate, will not be used if auto_lr_find is used.
+        dataset (Optional[str]): Dataset to used
+        save_images (Optional[bool]): Boolean to decide whether to save images
+        save_path (Optional[str]): Path to save images
+        """
+
+        super().__init__()
+        self.latent_size = latent_size
+        self.hidden_size = hidden_size
+        if save_images:
+            self.save_path = f'{save_path}/{kwargs["model_type"]}_images/'
+        self.save_hyperparameters()
+        self.save_images = save_images
+        self.lr = lr
+        self.batch_size = batch_size
+        self.encoder = nn.Sequential(
+            Flatten(),
+            nn.Linear(784, 392), nn.BatchNorm1d(392), nn.LeakyReLU(0.1),
+            nn.Linear(392, 196), nn.BatchNorm1d(196), nn.LeakyReLU(0.1),
+            nn.Linear(196, 128), nn.BatchNorm1d(128), nn.LeakyReLU(0.1),
+            nn.Linear(128, latent_size)
+        )
+        self.hidden2mu = nn.Linear(latent_size, latent_size)
+        self.hidden2log_var = nn.Linear(latent_size, latent_size)
+        self.alpha = alpha
+        self.decoder = nn.Sequential(
+            nn.Linear(latent_size, 128), nn.BatchNorm1d(128), nn.LeakyReLU(0.1),
+            nn.Linear(128, 196), nn.BatchNorm1d(196), nn.LeakyReLU(0.1),
+            nn.Linear(196, 392), nn.BatchNorm1d(392), nn.LeakyReLU(0.1),
+            nn.Linear(392, 784),
+            Stack(1, 28, 28),
+            nn.Tanh()
+        )
+        self.height = kwargs.get("height")
+        self.width = kwargs.get("width")
+        self.data_transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Lambda(lambda x:2*x-1.)])
+        self.dataset = dataset
+
+    def encode(self, x):
+        hidden = self.encoder(x)
+        mu = self.hidden2mu(hidden)
+        log_var = self.hidden2log_var(hidden)
+        return mu, log_var
+
+    def decode(self, x):
+        x = self.decoder(x)
+        return x
+
+    def reparametrize(self, mu, log_var):
+        # Reparametrization Trick to allow gradients to backpropagate from the
+        # stochastic part of the model
+        sigma = torch.exp(0.5*log_var)
+        z = torch.randn_like(sigma)
+        return mu + sigma*z
+
+    def training_step(self, batch, batch_idx):
+        x, _ = batch
+        mu, log_var, x_out = self.forward(x)
+        kl_loss = (-0.5*(1+log_var - mu**2 -
+                         torch.exp(log_var)).sum(dim=1)).mean(dim=0)
+        recon_loss_criterion = nn.MSELoss()
+        recon_loss = recon_loss_criterion(x, x_out)
+        # print(kl_loss.item(),recon_loss.item())
+        loss = recon_loss*self.alpha + kl_loss
+
+        self.log('train_loss', loss, on_step=False,
+                 on_epoch=True, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, _ = batch
+        mu, log_var, x_out = self.forward(x)
+
+        kl_loss = (-0.5*(1+log_var - mu**2 -
+                         torch.exp(log_var)).sum(dim=1)).mean(dim=0)
+        recon_loss_criterion = nn.MSELoss()
+        recon_loss = recon_loss_criterion(x, x_out)
+        # print(kl_loss.item(),recon_loss.item())
+        loss = recon_loss*self.alpha + kl_loss
+        self.log('val_kl_loss', kl_loss, on_step=False, on_epoch=True)
+        self.log('val_recon_loss', recon_loss, on_step=False, on_epoch=True)
+        self.log('val_loss', loss, on_step=False, on_epoch=True)
+        # print(x.mean(),x_out.mean())
+        return x_out, loss
+
+    def validation_epoch_end(self, outputs):
+        if not self.save_images:
+            return
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+        choice = random.choice(outputs)
+        output_sample = choice[0]
+        output_sample = output_sample.reshape(-1, 1, self.width, self.height)
+        # output_sample = self.scale_image(output_sample)
+        save_image(
+            output_sample,
+            f"{self.save_path}/epoch_{self.current_epoch+1}.png",
+            # value_range=(-1, 1)
+        )
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=(self.lr or self.learning_rate))
+        lr_scheduler = ReduceLROnPlateau(optimizer,)
+        return {
+            "optimizer": optimizer, "lr_scheduler": lr_scheduler,
+            "monitor": "val_loss"
+        }
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        hidden = self.reparametrize(mu, log_var)
+        output = self.decode(hidden)
+        return mu, log_var, output
+
+    # Functions for dataloading
+    def train_dataloader(self):
+        if self.dataset == "mnist":
+            train_set = MNIST('data/', download=True,
+                              train=True, transform=self.data_transform)
+        elif self.dataset == "fashion-mnist":
+            train_set = FashionMNIST(
+                'data/', download=True, train=True,
+                transform=self.data_transform)
+        elif self.dataset == "celeba":
+            train_set = CelebA('data/', download=False, split="train", transform=self.data_transform)
+        return DataLoader(train_set, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        if self.dataset == "mnist":
+            val_set = MNIST('data/', download=True, train=False,
+                            transform=self.data_transform)
+        elif self.dataset == "fashion-mnist":
+            val_set = FashionMNIST(
+                'data/', download=True, train=False,
+                transform=self.data_transform)
+        elif self.dataset == "celeba":
+            val_set = CelebA('data/', download=False, split="valid", transform=self.data_transform)
+        return DataLoader(val_set, batch_size=self.batch_size)
+
+    def scale_image(self, img):
+        out = (img + 1) / 2
+        return out
+
+    def interpolate(self, x1, x2):
+
+        assert x1.shape == x2.shape, "Inputs must be of the same shape"
+        if x1.dim() == 3:
+            x1 = x1.unsqueeze(0)
+        if x2.dim() == 3:
+            x2 = x2.unsqueeze(0)
+        if self.training:
+            raise Exception(
+                "This function should not be called when model is still "
+                "in training mode. Use model.eval() before calling the "
+                "function")
+        mu1, lv1 = self.encode(x1)
+        mu2, lv2 = self.encode(x2)
+        z1 = self.reparametrize(mu1, lv1)
+        z2 = self.reparametrize(mu2, lv2)
+        weights = torch.arange(0.1, 0.9, 0.1)
+        intermediate = [self.decode(z1)]
+        for wt in weights:
+            inter = (1.-wt)*z1 + wt*z2
+            intermediate.append(self.decode(inter))
+        intermediate.append(self.decode(z2))
+        out = torch.stack(intermediate, dim=0).squeeze(1)
+        return out, (mu1, lv1), (mu2, lv2)

requirements.txt

@@ -0,0 +1,8 @@
+torch==1.8.0
+python-box==5.3.0
+tensorboardX==2.1
+pydantic
+streamlit==0.82.0
+streamlit-drawable-canvas==0.8.0
+pytorch_lightning>=1.1.1
+torchvision==0.9.0

test.py

@@ -0,0 +1,22 @@
+from pytorch_lightning import Trainer
+from models import vae_models
+from config import config
+from pytorch_lightning.loggers import TensorBoardLogger
+import os
+
+def make_model(config):
+    model_type = config.model_type
+    model_config = config.model_config
+
+    if model_type not in vae_models.keys():
+        raise NotImplementedError("Model Architecture not implemented")
+    else:
+        return vae_models[model_type](**model_config.dict())
+
+
+if __name__ == "__main__":
+    model_type = config.model_type
+    model = vae_models[model_type].load_from_checkpoint("./saved_models/vae_alpha_1024_dim_128.ckpt")
+    logger = TensorBoardLogger(**config.log_config.dict())
+    trainer = Trainer(gpus=1, logger=logger)
+    trainer.test(model)

train.py

@@ -0,0 +1,37 @@
+from pytorch_lightning import Trainer
+from models import vae_models
+from config import config
+from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.loggers import TensorBoardLogger
+import os
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+
+
+def make_model(config):
+    model_type = config.model_type
+    model_config = config.model_config
+
+    if model_type not in vae_models.keys():
+        raise NotImplementedError("Model Architecture not implemented")
+    else:
+        return vae_models[model_type](**model_config.dict())
+
+
+if __name__ == "__main__":
+    model = make_model(config)
+    train_config = config.train_config
+    logger = TensorBoardLogger(**config.log_config.dict())
+    trainer = Trainer(**train_config.dict(), logger=logger,
+                      callbacks=LearningRateMonitor())
+    if train_config.auto_lr_find:
+        lr_finder = trainer.tuner.lr_find(model)
+        new_lr = lr_finder.suggestion()
+        print("Learning Rate Chosen:", new_lr)
+        model.lr = new_lr
+        trainer.fit(model)
+    else:
+        trainer.fit(model)
+    if not os.path.isdir("./saved_models"):
+        os.mkdir("./saved_models")
+    trainer.save_checkpoint(
+        f"saved_models/{config.model_type}_alpha_{config.model_config.alpha}_dim_{config.model_config.hidden_size}.ckpt")

utils.py

@@ -0,0 +1,63 @@
+from pytorch_lightning import Trainer
+from torchvision.utils import save_image
+from models import vae_models
+from config import config
+from PIL import Image 
+from pytorch_lightning.loggers import TensorBoardLogger
+import torch
+from torch.nn.functional import interpolate
+from torchvision.transforms import Resize, ToPILImage, Compose
+from torchvision.utils import make_grid
+
+def load_model(ckpt, model_type="vae"):
+    model = vae_models[model_type].load_from_checkpoint(f"./saved_models/{ckpt}")
+    model.eval()
+    return model
+
+def parse_model_file_name(file_name):
+    # Hard Coded Parsing based on the filenames that I use
+    substrings = file_name.split(".")[0].split("_")
+    name, alpha, dim = substrings[0], substrings[2], substrings[4]
+    new_name = ""
+    if name == "vae":
+        new_name += "Vanilla VAE"
+    new_name += f" | alpha={alpha}"
+    new_name += f" | dim={dim}"
+    return new_name
+
+def tensor_to_img(tsr):
+    if tsr.ndim == 4:
+        tsr = tsr.squeeze(0)
+    
+    transform = Compose([
+        ToPILImage()
+    ])
+    img = transform(tsr)
+    return img
+
+
+def resize_img(img, w, h):
+    return img.resize((w, h))
+
+def canvas_to_tensor(canvas):
+    """
+    Convert Image of RGBA to single channel B/W and convert from numpy array
+    to a PyTorch Tensor of [1,1,28,28]
+    """
+    img = canvas.image_data
+    img = img[:, :, :-1]  # Ignore alpha channel
+    img = img.mean(axis=2)
+    img = img/255
+    img = img*2 - 1.
+    img = torch.FloatTensor(img)
+    tens = img.unsqueeze(0).unsqueeze(0)
+    tens = interpolate(tens, (28, 28))
+    return tens
+
+
+def export_to_onnx(ckpt):
+    model = load_model(ckpt)
+    filepath = "model.onnx"
+    test_iter = iter(model.test_dataloader())
+    sample, _ = next(test_iter)
+    model.to_onnx(filepath, sample, export_params=True)