vae.py

# -*- coding: utf-8 -*-
"""VAE.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1qNjaOTR1h0A8jn2Cqs9yipIbOuQ6wsRR
"""

!pip install pytorch_lightning

import torch
import numpy as np
import torch.nn as nn
import pytorch_lightning as pl
from torch.optim import Adam
import matplotlib.pyplot as plt
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from mpl_toolkits.axes_grid1 import ImageGrid
from torchvision.utils import save_image, make_grid

class VAE(pl.LightningModule):

    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=200):
        super(VAE, self).__init__()

        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, latent_dim),
            nn.LeakyReLU(0.2)
        )

        self.mean_layer = nn.Linear(latent_dim, 2)
        self.logvar_layer = nn.Linear(latent_dim, 2)

        self.decoder = nn.Sequential(
            nn.Linear(2, latent_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(latent_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()
        )

    def encode(self, x):
        x = self.encoder(x)
        mean, logvar = self.mean_layer(x), self.logvar_layer(x)
        return mean, logvar

    def reparameterization(self, mean, var):
        epsilon = torch.randn_like(var)
        z = mean + var * epsilon
        return z

    def decode(self, x):
        return self.decoder(x)

    def forward(self, x):
        mean, logvar = self.encode(x)
        z = self.reparameterization(mean, torch.exp(0.5 * logvar))
        x_hat = self.decode(z)
        return x_hat, mean, logvar

    def training_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        x_hat, mean, logvar = self(x)
        loss = self.loss_function(x, x_hat, mean, logvar)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=1e-3)

    def loss_function(self, x, x_hat, mean, logvar):
        reproduction_loss = nn.functional.binary_cross_entropy(x_hat, x, reduction='sum')
        KLD = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp())
        return reproduction_loss + KLD

    def generate_digit(self,mean, var):
        z_sample = torch.tensor([[mean, var]], dtype=torch.float)
        x_decoded = self.decode(z_sample)
        digit = x_decoded.detach().cpu().reshape(28, 28) # reshape vector to 2d array
        plt.title(f'[{mean},{var}]')
        plt.imshow(digit, cmap='gray')
        plt.axis('off')
        plt.show()

    def plot_latent_space(self, scale=5.0, n=25, digit_size=28, figsize=15):
        # display a n*n 2D manifold of digits
        figure = np.zeros((digit_size * n, digit_size * n))

        # construct a grid
        grid_x = np.linspace(-scale, scale, n)
        grid_y = np.linspace(-scale, scale, n)[::-1]

        for i, yi in enumerate(grid_y):
            for j, xi in enumerate(grid_x):
                z_sample = torch.tensor([[xi, yi]], dtype=torch.float)
                x_decoded = self.decode(z_sample)
                digit = x_decoded[0].detach().cpu().reshape(digit_size, digit_size)
                figure[i * digit_size : (i + 1) * digit_size, j * digit_size : (j + 1) * digit_size,] = digit

        plt.figure(figsize=(figsize, figsize))
        plt.title('VAE Latent Space Visualization')
        start_range = digit_size // 2
        end_range = n * digit_size + start_range
        pixel_range = np.arange(start_range, end_range, digit_size)
        sample_range_x = np.round(grid_x, 1)
        sample_range_y = np.round(grid_y, 1)
        plt.xticks(pixel_range, sample_range_x)
        plt.yticks(pixel_range, sample_range_y)
        plt.xlabel("mean, z [0]")
        plt.ylabel("var, z [1]")
        plt.imshow(figure, cmap="Greys_r")
        plt.show()

class MNISTDataModule(pl.LightningDataModule):

    def __init__(self, path='~/datasets', batch_size=100):
        super().__init__()
        self.path = path
        self.batch_size = batch_size

    def prepare_data(self):
        transform = transforms.Compose([transforms.ToTensor()])
        MNIST(self.path, train=True, download=True, transform=transform)
        MNIST(self.path, train=False, download=True, transform=transform)

    def setup(self, stage=None):
        transform = transforms.Compose([transforms.ToTensor()])
        self.train_dataset = MNIST(self.path, train=True, transform=transform)
        self.test_dataset = MNIST(self.path, train=False, transform=transform)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False)

def main():
    data_module = MNISTDataModule()
    vae_model = VAE()
    trainer = pl.Trainer(max_epochs=3)
    trainer.fit(vae_model, datamodule=data_module)
    vae_model.generate_digit(0.0, 1.0), vae_model.generate_digit(1.0, 0.0)
    vae_model.plot_latent_space(scale=1.0)
if __name__ == "__main__":
    main()