Creating a Variational Autoencoder from Scratch Using PyTorch

A variational autoencoder (VAE) is a type of neural network that can learn to generate new data samples by training on a dataset. VAEs are a type of generative model that learn a probability distribution over the data, allowing them to generate new samples that are similar to those in the training data. In this tutorial, we will walk through how to build a VAE from scratch using PyTorch.

First, let’s define the architecture of the VAE. The VAE is composed of an encoder and a decoder. The encoder takes an input data sample and encodes it into a latent space representation, while the decoder takes a sample from the latent space and reconstructs the original input data. The VAE is trained by maximizing the evidence lower bound (ELBO), which consists of two parts: the reconstruction loss and the KL divergence between the latent distribution and a standard normal distribution.

Let’s start by defining the encoder and decoder networks in PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F

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_mean = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        h = F.relu(self.fc1(x))
        mean = self.fc_mean(h)
        logvar = self.fc_logvar(h)

        return mean, logvar

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 = F.relu(self.fc1(z))
        x_recon = torch.sigmoid(self.fc2(h))

        return x_recon

Next, we will define the VAE model that combines the encoder and decoder:

class VAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(VAE, self).__init__()

        self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
        self.decoder = Decoder(latent_dim, hidden_dim, input_dim)

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

    def forward(self, x):
        mean, logvar = self.encoder(x)
        z = self.reparameterize(mean, logvar)
        x_recon = self.decoder(z)

        return x_recon, mean, logvar

Now, we will define the loss function for training the VAE:

def loss_function(x_recon, x, mean, logvar):
    recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum')
    kl_div = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp())

    return recon_loss + kl_div

Finally, we can train the VAE on a dataset. Below is an example of training the VAE model on the MNIST dataset:

# Load the MNIST dataset
from torchvision import datasets, transforms

transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

# Initialize the VAE model
input_dim = 784
hidden_dim = 256
latent_dim = 20
vae = VAE(input_dim, hidden_dim, latent_dim)

# Define the optimizer
optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)

# Train the VAE
num_epochs = 10
for epoch in range(num_epochs):
    for i, (x, _) in enumerate(train_loader):
        x = x.view(-1, 784)
        x_recon, mean, logvar = vae(x)
        loss = loss_function(x_recon, x, mean, logvar)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
                epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

This is a basic implementation of a VAE in PyTorch. You can further customize the VAE by modifying the network architecture, loss function, or training parameters. VAEs are a powerful tool for generating new data samples and can be applied to various types of datasets.