🎓 A1Learner
🔍

PyTorch Training a VAE

Training a VAE means training the encoder, the sampler, and the decoder together on the two-part loss. This lesson trains a VAE and a plain autoencoder on the same blobs and prints both reconstruction errors.

Every rebuild samples a fresh code

Two draws from the same distribution give two different codes. The seed changes between the draws here on purpose, so the difference shows.

dist = torch.distributions.Normal(torch.tensor([0.0]), torch.tensor([1.0]))
torch.manual_seed(1)
print(round(dist.rsample().item(), 4))
torch.manual_seed(2)
print(round(dist.rsample().item(), 4))

A far mean costs more KL

The KL term grows fast as the mean moves away from 0. A mean of 2 costs sixteen times what a mean of one half costs.

for m in [0.5, 2.0]:
    mean = torch.tensor([m])
    logvar = torch.tensor([0.0])
    kl = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp())
    print(m, kl.item())

The two-part loss in the loop

The loop adds the reconstruction term and the KL term into one loss. One backward call then trains the whole model.

rebuilt, mean, logvar = vae(points)
recon = nn.functional.mse_loss(rebuilt, points)
kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
loss = recon + kl

Example

import torch
from torch import nn

torch.manual_seed(0)
per_blob = 50
blob_a = torch.randn(per_blob, 2) * 0.5 + torch.tensor([-2.0, -2.0])
blob_b = torch.randn(per_blob, 2) * 0.5 + torch.tensor([2.0, 2.0])
points = torch.cat([blob_a, blob_b])


class Autoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
        self.decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))

    def forward(self, x):
        return self.decoder(self.encoder(x))


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(nn.Linear(2, 16), nn.ReLU())
        self.to_mean = nn.Linear(16, 1)
        self.to_logvar = nn.Linear(16, 1)
        self.decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))

    def encode(self, x):
        h = self.body(x)
        return self.to_mean(h), self.to_logvar(h)

    def forward(self, x):
        mean, logvar = self.encode(x)
        std = torch.exp(0.5 * logvar)
        z = torch.distributions.Normal(mean, std).rsample()
        return self.decoder(z), mean, logvar


torch.manual_seed(1)
ae = Autoencoder()
opt = torch.optim.Adam(ae.parameters(), lr=0.02)
for epoch in range(300):
    loss = nn.functional.mse_loss(ae(points), points)
    opt.zero_grad()
    loss.backward()
    opt.step()

torch.manual_seed(1)
vae = VAE()
opt = torch.optim.Adam(vae.parameters(), lr=0.02)
for epoch in range(300):
    rebuilt, mean, logvar = vae(points)
    recon = nn.functional.mse_loss(rebuilt, points)
    kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
    loss = recon + kl
    opt.zero_grad()
    loss.backward()
    opt.step()

with torch.no_grad():
    ae_error = nn.functional.mse_loss(ae(points), points).item()
    torch.manual_seed(5)
    rebuilt, mean, logvar = vae(points)
    vae_error = nn.functional.mse_loss(rebuilt, points).item()
    blob_a_mean = mean[:per_blob].mean().item()
    blob_b_mean = mean[per_blob:].mean().item()

print("Autoencoder reconstruction error:", round(ae_error, 4))
print("VAE reconstruction error:", round(vae_error, 4))
print("Blob A average latent mean:", round(blob_a_mean, 2))
print("Blob B average latent mean:", round(blob_b_mean, 2))

Key points

  • Both models here train on the same 100 points for 300 passes
  • The VAE pays for its ordered latent space with a higher reconstruction error
  • The two blobs take opposite ends of the latent line after training
  • A latent mean is the average code the encoder gives one blob

In this example, the plain autoencoder from the earlier lesson trains first on the 100 blob points. Its reconstruction error prints as 0.0876. The VAE then trains on the same points with the same seed and the same 300 passes. Its reconstruction error prints as 0.459, about five times higher. The VAE does not win on reconstruction, and the numbers say so plainly. Its gain shows in the latent means. Blob A averages -0.9 on the latent line, and blob B averages 0.88. The KL term forced the codes into one ordered line, and the blobs took opposite ends of it.

Note

The VAE error is measured on a seeded sample, so the number repeats on every run

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A trained VAE rebuilds worse than a plain autoencoder on the same data. What does the VAE gain for that price?