Variational Autoencoder (VAE)

import torch, torch.nn as nn
from sklearn.datasets import load_digits

torch.manual_seed(0)
X = torch.tensor(load_digits().data / 16.0, dtype=torch.float32)   # 8x8 digits

class VAE(nn.Module):
    def __init__(self, z=2):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 2 * z))
        self.dec = nn.Sequential(nn.Linear(z, 128), nn.ReLU(), nn.Linear(128, 64), nn.Sigmoid())

    def forward(self, x):
        mu, logvar = self.enc(x).chunk(2, dim=1)
        z = mu + torch.randn_like(mu) * torch.exp(0.5 * logvar)    # reparameterization trick
        return self.dec(z), mu, logvar

vae = VAE()
opt = torch.optim.Adam(vae.parameters(), lr=1e-3)
for epoch in range(200):
    recon, mu, logvar = vae(X)
    rec = nn.functional.binary_cross_entropy(recon, X, reduction="sum") / len(X)
    kl = -0.5 * torch.sum(1 + logvar - mu ** 2 - logvar.exp()) / len(X)
    loss = rec + kl
    opt.zero_grad(); loss.backward(); opt.step()
    if epoch % 50 == 49:
        print(f"epoch {epoch + 1}  reconstruction {rec.item():.1f}  KL {kl.item():.2f}")

# Generate: decode random points from N(0, I)
with torch.no_grad():
    sample = vae.dec(torch.randn(1, 2)).view(8, 8)
for row in sample:
    print("".join(" .:-=+*#%@"[int(v * 9.99)] for v in row))