Implementing a Variational Autoencoder for Handwritten Digits in PyTorch: A Code Example

In this tutorial, we will explore a Variational Autoencoder (VAE) for generating handwritten digits using PyTorch. VAEs are a type of generative model that learn to encode and decode data in a continuous latent space. They are particularly useful for generating new samples from a dataset and are widely used in fields such as image generation and machine learning.

We will use the MNIST dataset, which consists of 28×28 pixel grayscale images of handwritten digits ranging from 0 to 9. We will build a VAE model that can generate new handwritten digits based on this dataset.

First, let’s install the necessary libraries:

pip install torch torchvision matplotlib

Next, let’s import the required libraries in our Python script:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import transforms
import matplotlib.pyplot as plt

Now, let’s define the VAE model. The model consists of an encoder network, a decoder network, and a reparameterization layer:

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

        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU()
        )

        self.mu = nn.Linear(256, latent_dim)
        self.log_var = nn.Linear(256, latent_dim)

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, input_dim),
            nn.Sigmoid()
        )

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

    def forward(self, x):
        x = self.encoder(x)
        mu = self.mu(x)
        log_var = self.log_var(x)

        z = self.reparameterize(mu, log_var)
        x_reconstructed = self.decoder(z)

        return x_reconstructed, mu, log_var

Next, let’s define the loss function for the VAE, which consists of a reconstruction loss and a KL divergence loss:

def vae_loss(x, x_reconstructed, mu, log_var):
    # Reconstruction Loss
    reconstruction_loss = F.binary_cross_entropy(x_reconstructed, x, reduction='sum')

    # KL Divergence Loss
    kl_divergence_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

    return reconstruction_loss + kl_divergence_loss

Now, let’s define the training loop for the VAE model:

def train_vae(model, dataloader, optimizer, num_epochs):
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0.0

        for i, (x, _) in enumerate(dataloader):
            x = x.view(-1, 28*28)
            x_reconstructed, mu, log_var = model(x)

            loss = vae_loss(x, x_reconstructed, mu, log_var)
            total_loss += loss.item()

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

        print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataloader.dataset)}')

Finally, let’s load the MNIST dataset, initialize the VAE model, and train the model:

# Load MNIST dataset
transform = transforms.Compose([
    transforms.ToTensor()
])

mnist_dataset = MNIST(root='./data', train=True, download=True, transform=transform)
dataloader = DataLoader(mnist_dataset, batch_size=128, shuffle=True)

# Initialize VAE model
input_dim = 28*28
latent_dim = 64
model = VAE(input_dim, latent_dim)

# Set optimizer and number of epochs
optimizer = optim.Adam(model.parameters(), lr=1e-3)
num_epochs = 10

# Train the model
train_vae(model, dataloader, optimizer, num_epochs)

After training the VAE model, we can generate new handwritten digits by sampling from the latent space:

def generate_samples(model, num_samples):
    model.eval()

    with torch.no_grad():
        z = torch.randn(num_samples, latent_dim)
        generated_samples = model.decoder(z)

    return generated_samples

# Generate new samples
num_samples = 10
generated_samples = generate_samples(model, num_samples)

# Display generated samples
fig, axs = plt.subplots(1, num_samples, figsize=(20, 4))
for i in range(num_samples):
    axs[i].imshow(generated_samples[i].view(28, 28).numpy(), cmap='gray')
    axs[i].axis('off')

plt.show()

And that’s it! We have successfully built a Variational Autoencoder for generating handwritten digits using PyTorch. Feel free to experiment with different hyperparameters, network architectures, and datasets to further improve the model’s performance.