vae_mlp_mnist.py

"""
Install pytorch lightning and einops

pip install pytorch_lightning einops
"""

import torch
from torch import nn
import pytorch_lightning as pl
from torch.nn import functional as F
from torchvision.datasets import MNIST
from torchvision import transforms
from argparse import ArgumentParser
from torch.utils.data import DataLoader
from einops import rearrange


class VAE(nn.Module):
    def __init__(self, n_z, model_name="vae"):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 512), nn.ReLU())
        self.model_name = model_name
        self.fc_mu = nn.Linear(512, n_z)
        self.fc_var = nn.Linear(512, n_z)
        self.decoder = nn.Sequential(nn.Linear(n_z, 512), nn.ReLU(), nn.Linear(512, 28 * 28), nn.Sigmoid())

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        x = self.encoder(x)
        mu = self.fc_mu(x)
        log_var = self.fc_var(x)
        p, q, z = self.sample(mu, log_var)
        return self.decoder(z)

    def _run_step(self, x):
        x = self.encoder(x)
        mu = self.fc_mu(x)
        log_var = self.fc_var(x)
        p, q, z = self.sample(mu, log_var)
        return z, self.decoder(z), p, q

    def encode(self, x):
        x = self.encoder(x)
        mu = self.fc_mu(x)
        return mu

    def sample(self, mu, log_var):
        std = torch.exp(log_var / 2)
        p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
        q = torch.distributions.Normal(mu, std)
        z = q.rsample()
        return p, q, z


class BasicVAEModule(pl.LightningModule):
    def __init__(self, n_z=2, kl_coeff=0.1, lr=0.001):
        super().__init__()
        self.vae = VAE(n_z)
        self.kl_coeff = kl_coeff
        self.lr = lr

    def forward(self, x):
        return self.vae(x)

    def step(self, batch, batch_idx):
        x, y = batch
        z, x_hat, p, q = self.vae._run_step(x)

        recon_loss = F.binary_cross_entropy(x_hat, x, reduction="sum")

        log_qz = q.log_prob(z)
        log_pz = p.log_prob(z)

        kl = log_qz - log_pz
        kl = kl.sum()  # I tried sum, here
        kl *= self.kl_coeff

        loss = kl + recon_loss

        logs = {
            "recon_loss": recon_loss,
            "kl": kl,
            "loss": loss,
        }
        return loss, logs

    def training_step(self, batch, batch_idx):
        loss, logs = self.step(batch, batch_idx)
        self.log_dict({f"train_{k}": v for k, v in logs.items()}, on_step=True, on_epoch=False)
        return loss

    def validation_step(self, batch, batch_idx):
        loss, logs = self.step(batch, batch_idx)
        self.log_dict({f"val_{k}": v for k, v in logs.items()})
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)


if __name__ == "__main__":
    parser = ArgumentParser(description="Hyperparameters for our experiments")
    parser.add_argument("--latent-dim", type=int, default=12, help="size of latent dim for our vae")
    parser.add_argument("--epochs", type=int, default=50, help="num epochs")
    parser.add_argument("--gpus", type=int, default=1, help="gpus, if no gpu set to 0, to run on all  gpus set to -1")
    parser.add_argument("--bs", type=int, default=500, help="batch size")
    hparams = parser.parse_args()

    mnist_full = MNIST(
        ".",
        download=True,
        train=True,
        transform=transforms.Compose([transforms.ToTensor(), lambda x: rearrange(x, "c h w -> (c h w)")]),
    )
    dm = DataLoader(mnist_full, batch_size=hparams.bs, shuffle=True)
    vae = BasicVAEModule(hparams.latent_dim)

    trainer = pl.Trainer(gpus=hparams.gpus, weights_summary="full", max_epochs=hparams.epochs)
    trainer.fit(vae, dm)
    torch.save(vae.state_dict(), "vae-mnist-mlp.ckpt")