变分自编码器 (VAE)：从直觉到实现与调试

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def reparameterize(mu, logvar):
    """以可微的方式从 N(mu, sigma^2) 采样 z。

    Args:
        mu:     (B, latent_dim) 编码器均值
        logvar: (B, latent_dim) 编码器对数方差
    Returns:
        z:      (B, latent_dim) 采样结果
    """
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return mu + std * eps

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

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

class Encoder(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__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 = F.relu(self.fc1(x))
        return self.fc_mu(h), self.fc_logvar(h)

class Decoder(nn.Module):
    def __init__(self, latent_dim=20, hidden_dim=400, output_dim=784):
        super().__init__()
        self.fc3 = nn.Linear(latent_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, output_dim)

    def forward(self, z):
        h = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h))   # 像素概率，落在 [0, 1]

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
        self.decoder = Decoder(latent_dim, hidden_dim, input_dim)

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

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

def vae_loss(recon_x, x, mu, logvar, beta=1.0):
    """负 ELBO = 重建项 + beta * KL。

    返回的是 batch 上的**总和**而不是均值，因为 KL 闭式解本身也是求和。
    日志输出时按需除以 batch_size 即可换成"每张图"的数值。
    """
    # 伯努利重建（适用于二值化 MNIST）；
    # 连续值数据用 F.mse_loss(reduction='sum')。
    recon = F.binary_cross_entropy(recon_x, x, reduction="sum")

    # KL(N(mu, sigma^2) || N(0, I))，对 batch 与潜在维度求和。
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return recon + beta * kl, recon, kl

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

def train_vae(model, loader, optimizer, epochs=20, device="cuda",
              warmup_epochs=10):
    model.to(device).train()
    for epoch in range(1, epochs + 1):
        # KL 退火：beta 在前 warmup_epochs 轮里从 0 线性升到 1。
        beta = min(1.0, epoch / warmup_epochs)
        ep_total = ep_recon = ep_kl = 0.0
        for x, _ in loader:
            x = x.view(-1, 784).to(device)
            recon_x, mu, logvar = model(x)
            loss, recon, kl = vae_loss(recon_x, x, mu, logvar, beta=beta)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            ep_total += loss.item()
            ep_recon += recon.item()
            ep_kl    += kl.item()

        n = len(loader.dataset)
        print(f"epoch {epoch:3d}  beta={beta:.2f}  "
              f"loss={ep_total/n:.2f}  recon={ep_recon/n:.2f}  "
              f"kl={ep_kl/n:.2f}")

1
2
3

logvar = torch.clamp(logvar, min=-10, max=10)        # 给 exp(logvar) 兜底
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
# 优先 AdamW，lr 取 1e-3 ~ 3e-4，weight_decay 取 1e-5 量级

1
2
3
4
5
6
7
8

# MNIST：ToTensor() 已经给到 [0, 1]
transform = transforms.ToTensor()

# 连续数据：标准化到 [-1, 1]
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5]),
])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import matplotlib.pyplot as plt

@torch.no_grad()
def plot_latent_space(model, loader, device="cuda"):
    model.eval()
    mus, labels = [], []
    for x, y in loader:
        mu, _ = model.encoder(x.view(-1, 784).to(device))
        mus.append(mu.cpu()); labels.append(y)
    mus    = torch.cat(mus).numpy()
    labels = torch.cat(labels).numpy()
    plt.scatter(mus[:, 0], mus[:, 1], c=labels, cmap="tab10",
                s=8, alpha=0.7)
    plt.colorbar(); plt.title("VAE 潜在空间 (2D)"); plt.show()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

@torch.no_grad()
def sample_vae(model, n=16, device="cuda"):
    model.eval()
    z = torch.randn(n, model.encoder.fc_mu.out_features).to(device)
    return model.decoder(z).cpu().view(-1, 28, 28)

@torch.no_grad()
def interpolate(model, x1, x2, steps=10, device="cuda"):
    model.eval()
    mu1, _ = model.encoder(x1.view(1, -1).to(device))
    mu2, _ = model.encoder(x2.view(1, -1).to(device))
    ts = torch.linspace(0, 1, steps + 1, device=device).view(-1, 1)
    z  = (1 - ts) * mu1 + ts * mu2
    return model.decoder(z).cpu().view(-1, 28, 28)