Building a Variational Autoencoder in PyTorch from Scratch

Variational Autoencoders (VAEs) are a type of generative model that have gained popularity due to their ability to generate new samples from a learned distribution. They offer a more elegant way of capturing the underlying distribution of data compared to traditional autoencoders because they learn a probability density over the set of inputs, rather than mapping each input to a single point in a latent space. In this article, we'll walk through building a VAE using PyTorch from scratch.

Prerequisites

Before diving into code, ensure you have:

• Familiarity with PyTorch basics.
• Understanding of neural networks and backpropagation.
• Basic knowledge of probability and statistics.

Setting Up the Environment

Let's start by installing PyTorch if you haven't already. You can install it using pip:

pip install torch torchvision

Now, import the necessary libraries:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

Define the Model Architecture

A VAE consists of two primary components: an encoder and a decoder. The encoder learns a representation (mean and variance), while the decoder generates data from this representation. Let's break it down:

Encoder

class Encoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(Encoder, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        h = torch.relu(self.fc1(x))
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar

Decoder

class Decoder(nn.Module):
    def __init__(self, latent_dim, hidden_dim, output_dim):
        super(Decoder, self).__init__()
        self.fc1 = nn.Linear(latent_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, z):
        h = torch.relu(self.fc1(z))
        return torch.sigmoid(self.fc2(h))

VAE Class

The VAE class combines both the encoder and decoder. It also handles sampling from the latent space:

class VAE(nn.Module):
    def __init__(self, encoder, decoder):
        super(VAE, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encoder(x)
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar

Training the VAE

Next, we need a way to train our model. We'll use the loss function that combines the reconstruction loss and the KL-divergence, which regularizes our learned distribution:

def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.binary_cross_entropy(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

Now set up the training loop:

epochs = 10
batch_size = 64
learning_rate = 1e-3

train_loader = DataLoader(..., batch_size=batch_size, shuffle=True)

vae = VAE(Encoder(784, 400, 20), Decoder(20, 400, 784))
optimizer = optim.Adam(vae.parameters(), lr=learning_rate)

for epoch in range(epochs):
    for data, _ in train_loader:
        data = data.view(-1, 784)  # Flatten the images
        optimizer.zero_grad()
        recon, mu, logvar = vae(data)
        loss = loss_function(recon, data, mu, logvar)
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

Conclusion

Building a VAE in PyTorch allows you to delve deeply into understanding more about deep learning models and their architectures. It's a flexible and powerful framework to create generative models, well-suited for many deep learning tasks. Once you master the basics of VAEs, you can further explore enhancements like conditioning, disentanglements, and more.