""" vae_lab.py —— 第 9 讲的工具箱：普通自编码器、变分自编码器（VAE）、VQ-VAE 与它们的训练 ===================================================================== 【这个文件要回答的问题】 slides 推导 VAE 的六步： 1. 自编码器不能采样 → 定义一个能采样的解码器：p(z) = N(0, I)，p_θ(x | z) = N(μ_θ(z), Σ_θ(z))； 2. 最大似然 log p_θ(x) = log ∫ p_θ(x | z) p(z) dz 算不出来 → 引入编码器 q_φ(z | x)，改为最大化下界 ELBO； 3. 交替优化 φ 和 θ 太慢 → 用随机梯度同时更新两者； 4. ELBO 的第一项（期望）算不出来 → 蒙特卡洛估计； 5. 直接对采样求梯度方差太大（得分函数估计）→ 重参数化：z = μ_φ(x) + σ_φ(x) ⊙ ε； 6. 实现：每张训练图取 S = 1 个样本，损失 = −log p_θ(x | z) + KL(q_φ(z | x) ‖ p(z))，反向传播，梯度下降。 这个文件把这六步写成代码，再加上 VQ-VAE（离散的隐变量 + 码本）。 【本文件提供】 AE 普通自编码器（slides："不是变分的那种"）：z = ENCODER(x)，x' = DECODER(z)，最小化 ‖x − x'‖² VAE slides 的 VAE：编码器、解码器都是一层隐藏层的 MLP（h = tanh(W z + b)，和 Kingma & Welling 2014 一样）； 解码器是高斯 N(μ_θ(z), σ² I)：σ 可以固定（超参数），也可以像 slides 那样由网络输出（sigma=None） vae_terms 一张图的 −log p_θ(x | z) 和 KL(q ‖ p)（KL 用两个高斯之间的闭式解） train_ae / train_vae VQVAE / train_vqvae VQ-VAE（van den Oord et al., 2017）：卷积编码器 → 7×7 的向量 → 换成码本里最近的向量 → 卷积解码器 code_prior / sample_codes 编号的自回归先验（第 2 讲的 GPT）与采样 【约定】图像是 (N, 1, 28, 28)、像素在 [0, 1]；VAE 的 μ_θ(z) 经过 sigmoid 也在 [0, 1]。 KL、对数似然的单位都是 nat，按"每张图"求和（不是每个像素平均）——这样 ELBO 才是 log p(x) 的下界。 """ import math import os import sys import torch import torch.nn as nn import torch.nn.functional as F _HERE = os.path.dirname(os.path.abspath(__file__)) for _d in ("lecture02", "lecture01"): _p = os.path.join(_HERE, "..", _d) if _p not in sys.path: sys.path.insert(0, _p) from lm_models import GPT # noqa: E402 第 2 讲的 GPT：VQ-VAE 编号的自回归先验 LOG_2PI = math.log(2 * math.pi) # ===================================================================== # 1. 普通自编码器 # ===================================================================== class AE(nn.Module): """784 → 512（tanh）→ z_dim → 512（tanh）→ 784（sigmoid）。确定性的编码和解码。""" def __init__(self, z_dim=16, hidden=512): super().__init__() self.z_dim = z_dim self.enc = nn.Sequential(nn.Flatten(), nn.Linear(784, hidden), nn.Tanh(), nn.Linear(hidden, z_dim)) self.dec = nn.Sequential(nn.Linear(z_dim, hidden), nn.Tanh(), nn.Linear(hidden, 784)) def encode(self, x): return self.enc(x) def decode(self, z): return torch.sigmoid(self.dec(z)).view(-1, 1, 28, 28) def forward(self, x): return self.decode(self.encode(x)) def train_ae(model, x, steps, batch=128, lr=1e-3, seed=0): """最小化重建误差 ‖x − DECODER(ENCODER(x))‖²（每张图求和）。【返回】最后 200 步的平均损失""" gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) hist = [] for _ in range(steps): xb = x[torch.randint(len(x), (batch,), generator=gen)] loss = ((model(xb) - xb) ** 2).flatten(1).sum(1).mean() opt.zero_grad() loss.backward() opt.step() hist.append(loss.item()) model.eval() return sum(hist[-200:]) / 200 # ===================================================================== # 2. VAE # ===================================================================== class VAE(nn.Module): """ slides 的 VAE（也是 Kingma & Welling 2014 的结构）： 编码器 q_φ(z | x) = N(μ_φ(x), diag σ_φ²(x))： h = tanh(W x + b)，μ_φ = W_μ h + b_μ，log σ_φ² = W_v h + b_v 先验 p(z) = N(0, I) 解码器 p_θ(x | z) = N(μ_θ(z), σ² I)： h = tanh(W₁ z + b₁)，μ_θ = sigmoid(W₂ h + b₂) σ：sigma 是一个数时固定；sigma=None 时像 slides 一样由网络给出：log σ² = W₃ h + b₃（每张图一个数） 【参数】z_dim：隐变量维数；hidden：隐藏层宽度；sigma：解码器的标准差（None 表示学出来） 【为什么编码器输出 log σ² 而不是 σ】σ 必须为正，而网络的输出可以是任意实数；取指数之后自动为正。 """ def __init__(self, z_dim=16, hidden=512, sigma=0.3): super().__init__() self.z_dim, self.sigma = z_dim, sigma self.enc = nn.Sequential(nn.Flatten(), nn.Linear(784, hidden), nn.Tanh()) self.mu = nn.Linear(hidden, z_dim) self.logvar = nn.Linear(hidden, z_dim) self.dec = nn.Sequential(nn.Linear(z_dim, hidden), nn.Tanh()) self.out = nn.Linear(hidden, 784) self.out_logvar = nn.Linear(hidden, 1) if sigma is None else None def encode(self, x): """【返回】(μ_φ(x), log σ_φ²(x))，形状都是 (N, z_dim)""" h = self.enc(x) return self.mu(h), self.logvar(h) def decode(self, z): """【返回】(μ_θ(z) 形状 (N, 1, 28, 28), log σ²（形状 (N, 1, 1, 1)）)""" h = self.dec(z) mean = torch.sigmoid(self.out(h)).view(-1, 1, 28, 28) if self.out_logvar is None: lv = torch.full((len(z), 1, 1, 1), 2 * math.log(self.sigma)) else: lv = self.out_logvar(h).view(-1, 1, 1, 1) return mean, lv def sample(self, n, gen): """从先验采样并解码；返回解码器的均值 μ_θ(z)（通常展示的"VAE 样本"就是它，不再加解码器的噪声）。""" return self.decode(torch.randn(n, self.z_dim, generator=gen))[0] def reparameterize(mu, logvar, gen): """重参数化：z = μ + σ ⊙ ε，ε ~ N(0, I)。随机性全在 ε 里，μ 和 σ 是网络的确定性输出，梯度可以穿过它们。""" return mu + (0.5 * logvar).exp() * torch.randn(mu.shape, generator=gen) def gaussian_nll(x, mean, logvar): """−log N(x; mean, σ² I)，对一张图的所有像素求和。【返回】(N,)""" return 0.5 * (((x - mean) ** 2) / logvar.exp() + logvar + LOG_2PI).flatten(1).sum(1) def kl_to_standard_normal(mu, logvar): """KL(N(μ, diag σ²) ‖ N(0, I)) = ½ Σ (μ² + σ² − 1 − log σ²)。【返回】(N,)，以及每一维的值 (N, z_dim)""" per_dim = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar) return per_dim.sum(1), per_dim def vae_terms(model, x, gen): """ slides 第 6 步：取 S = 1 个样本，算 −log p_θ(x | z) 和 KL(q_φ(z | x) ‖ p(z))。 【返回】(重建项 (N,), KL (N,), 每一维的 KL (N, z_dim))；负 ELBO = 重建项 + KL """ mu, logvar = model.encode(x) z = reparameterize(mu, logvar, gen) mean, out_lv = model.decode(z) rec = gaussian_nll(x, mean, out_lv) kl, per_dim = kl_to_standard_normal(mu, logvar) return rec, kl, per_dim def train_vae(model, x, steps, batch=128, lr=1e-3, seed=0): """ 最小化负 ELBO（重建项 + KL），编码器和解码器的参数同时更新（slides 第 3 步）。 【返回】最后 200 步的 (平均重建项, 平均 KL) """ gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) hist = [] for _ in range(steps): xb = x[torch.randint(len(x), (batch,), generator=gen)] rec, kl, _ = vae_terms(model, xb, gen) loss = (rec + kl).mean() opt.zero_grad() loss.backward() opt.step() hist.append((rec.mean().item(), kl.mean().item())) model.eval() tail = hist[-200:] return sum(h[0] for h in tail) / len(tail), sum(h[1] for h in tail) / len(tail) # ===================================================================== # 3. VQ-VAE # ===================================================================== class VQVAE(nn.Module): """ VQ-VAE：编码器输出一个 7×7 的"向量网格"，每个向量换成码本（K 个可学习的向量）里离它最近的那个，再交给解码器。 于是一张图被表示成 49 个**离散的**编号（每个在 0..K−1 之间）——和一句话被表示成 49 个词元一样。 编码器：像素重排 (1, 28, 28) → (4, 14, 14) → 3×3 卷积 → ReLU → 步长 2 的卷积 → (c, 7, 7) → ReLU → 1×1 卷积 → (D, 7, 7) 量化： 每个位置的 D 维向量 z_e → 码本里最近的 e_k（编号 k） 解码器：1×1 卷积 → ReLU → 转置卷积 → (c, 14, 14) → ReLU → 3×3 卷积 → (4, 14, 14) → 像素重排的逆 → sigmoid 【不可导的"取最近"怎么训练】直通估计（straight-through）：前向用 z_q = e_k，反向传播时假装 z_q = z_e， 把解码器对 z_q 的梯度原封不动地交给编码器。代码里写成 z_q = z_e + (e_k − z_e).detach()。 【码本怎么学】损失里加两项：‖sg[z_e] − e_k‖²（把码本向量拉向编码器的输出）+ β‖z_e − sg[e_k]‖²（"承诺"损失， 让编码器的输出别离码本太远）。sg = stop gradient（.detach()）。 """ def __init__(self, K=64, D=16, c=32): super().__init__() self.K, self.D = K, D self.enc = nn.Sequential(nn.Conv2d(4, c, 3, padding=1), nn.ReLU(), nn.Conv2d(c, c, 4, 2, 1), nn.ReLU(), nn.Conv2d(c, D, 1)) self.codebook = nn.Embedding(K, D) nn.init.uniform_(self.codebook.weight, -1.0 / K, 1.0 / K) self.dec = nn.Sequential(nn.Conv2d(D, c, 1), nn.ReLU(), nn.ConvTranspose2d(c, c, 4, 2, 1), nn.ReLU(), nn.Conv2d(c, 4, 3, padding=1)) def encode(self, x): """【返回】z_e，形状 (N, D, 7, 7)""" return self.enc(F.pixel_unshuffle(x, 2)) def quantize(self, z_e): """每个位置换成最近的码本向量。【返回】(z_q (N, D, 7, 7), 编号 (N, 7, 7))""" N, D, H, W = z_e.shape flat = z_e.permute(0, 2, 3, 1).reshape(-1, D) # (N·49, D) d2 = (flat ** 2).sum(1, keepdim=True) - 2 * flat @ self.codebook.weight.T + (self.codebook.weight ** 2).sum(1) idx = d2.argmin(1) z_q = self.codebook(idx).view(N, H, W, D).permute(0, 3, 1, 2) return z_q, idx.view(N, H, W) def decode_indices(self, idx): """编号 (N, 7, 7) → 图像 (N, 1, 28, 28)。""" z_q = self.codebook(idx).permute(0, 3, 1, 2) return self.decode(z_q) def decode(self, z_q): return torch.sigmoid(F.pixel_shuffle(self.dec(z_q), 2)) def forward(self, x, beta=0.25): """【返回】(重建图, 编号, 重建损失, 码本损失, 承诺损失)；损失都是对一批取平均（重建按每张图求和）""" z_e = self.encode(x) e, idx = self.quantize(z_e) z_q = z_e + (e - z_e).detach() # 直通估计 recon = self.decode(z_q) rec = ((recon - x) ** 2).flatten(1).sum(1).mean() codebook = ((z_e.detach() - e) ** 2).sum(1).mean() commit = ((z_e - e.detach()) ** 2).sum(1).mean() return recon, idx, rec, codebook, beta * commit def train_vqvae(model, x, steps, batch=64, lr=1e-3, seed=0, callback=None, every=0): """训练 VQ-VAE。【返回】最后 200 步的平均重建损失""" gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) hist = [] for step in range(1, steps + 1): xb = x[torch.randint(len(x), (batch,), generator=gen)] _, idx, rec, cb, com = model(xb) loss = rec + cb + com opt.zero_grad() loss.backward() opt.step() hist.append(rec.item()) if callback is not None and every and step % every == 0: callback(step, model, hist) model.eval() return sum(hist[-200:]) / len(hist[-200:]) @torch.no_grad() def code_usage(model, x, batch=500): """所有图的编号里，K 个码本向量各被用了多少次。【返回】(K,) 的计数，以及"困惑度"exp(熵)（有效使用的码本向量个数）""" counts = torch.zeros(model.K) for i in range(0, len(x), batch): _, idx = model.quantize(model.encode(x[i:i + batch])) counts += torch.bincount(idx.flatten(), minlength=model.K).float() p = counts / counts.sum() nz = p > 0 return counts, float(torch.exp(-(p[nz] * p[nz].log()).sum())) def code_prior(K=64): """ VQ-VAE 编号的自回归先验：第 2 讲的小 GPT（d = 64、2 层、4 个头）。 词表 = K 个编号 + 1 个开头符号 BOS（编号 K）；一张图 = [BOS, 49 个编号]（按行展开 7×7 的网格）。 """ return GPT(K + 1, d_model=64, n_layers=2, n_heads=4, max_len=50) @torch.no_grad() def sample_codes(gpt, n, gen, K=64, temperature=1.0): """从 GPT 先验一个一个地采 49 个编号（祖先采样，第 3 讲）。【返回】(n, 7, 7)""" gpt.eval() seq = torch.full((n, 1), K, dtype=torch.long) for _ in range(49): logits, _ = gpt(seq) probs = F.softmax(logits[:, -1, :K] / temperature, dim=-1) # 只在 K 个编号里采（不采 BOS） seq = torch.cat([seq, torch.multinomial(probs, 1, generator=gen)], dim=1) return seq[:, 1:].view(n, 7, 7) def count_parameters(model): return sum(p.numel() for p in model.parameters())