""" ============================================================================= CS236 · 第 6 讲共享模块：vae.py —— 像素世界上的变分自编码器（VAE） ============================================================================= 【这个模块是干什么的】 第 5 讲在"后验能精确算"的世界里推导了 ELBO。这里实现一个真正的深度隐变量模型——VAE： 先验 p(z) = N(0, I)，z ∈ R^{d_z} 解码器 p_θ(x | z) = Π_i Bernoulli(x_i; σ(f_θ(z)_i)) f_θ：两层 MLP，输出 64 个 logit 编码器 q_φ(z | x) = N(μ_φ(x), diag(σ_φ(x)²)) 一个 MLP 同时输出 μ 和 log σ² 训练目标是 ELBO（原课第 6 讲 slide 19–20）： L(x; θ, φ) = E_{q_φ(z|x)}[log p_θ(x | z)] − KL(q_φ(z | x) ‖ p(z)) 期望用**重参数化**的样本估计：z = μ + σ ⊙ ε，ε ~ N(0, I)；两个高斯之间的 KL 有闭式解。 主要接口： VAE(d_z, hidden) 模型 .encode(x) → (μ, log σ²) .decode_logits(z) → (…, 64) .elbo(x, generator, k=1, beta=1) 每张图的 ELBO 及两项：重构 rec、KL（单位 nat） .iwae(x, k, generator) 每张图的 IWAE 下界 L_k = E[log (1/k) Σ_j w_j]（单位 nat） .log_px(x, k, generator) 用 k 个重要性样本估计 log p_θ(x)（k 很大时就是"测试集似然"的常用估计） .sample(n, generator) z ~ N(0, I) → 解码 → 每个像素按伯努利采样 exact_log_px_2d(model, X) d_z = 2 时在 z 的细网格上数值积分，得到"精确的" log p_θ(x) train_vae(model, train, val, ...) Adam 训练（目标可以是 ELBO 或 IWAE，可选 β 和 KL 预热） get_vae(name, data) 有缓存（checkpoints/name.pt）就读，否则按 CONFIGS 训练并保存 load_l3() 导入第 3 讲的 ar_models（数据 pixel_data、真实熵 true_bits） 【运行方式】 python3 vae.py # 自检：KL 的闭式解 vs 蒙特卡洛；IWAE 随 k 单调；网格积分的收敛（约 10 秒，第一次会训练一个小模型） 【PyTorch 小知识】 - torch.randn(..., generator=g)：用自己的随机数生成器抽标准正态噪声，训练和评估都可以逐位复现。 - tensor.chunk(2, dim=-1)：沿最后一维切成两半（编码器的输出前一半是 μ，后一半是 log σ²）。 - 编码器输出 log σ² 而不是 σ：网络的输出可以是任意实数，exp 之后自动为正。 """ import math import os import sys import time import torch import torch.nn as nn import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) CKPT_DIR = os.path.join(HERE, "checkpoints") LN2 = math.log(2.0) LOG2PI = math.log(2 * math.pi) def load_l3(): """导入第 3 讲的 ar_models 模块（只读使用：pixel_data、true_bits、MADE 等），返回模块对象。""" sys.path.insert(0, os.path.join(HERE, "..", "lecture03")) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) import ar_models return ar_models def gaussian_log_prob(z, mu, logvar): """对角高斯 log N(z; μ, diag(exp(logvar)))，对最后一维求和。""" return -0.5 * (((z - mu) ** 2) / logvar.exp() + logvar + LOG2PI).sum(-1) def std_normal_log_prob(z): """标准正态 log N(z; 0, I)，对最后一维求和。""" return -0.5 * (z ** 2 + LOG2PI).sum(-1) class VAE(nn.Module): """ 伯努利解码器 + 对角高斯编码器的 VAE（见文件开头）。 【参数】d_z：隐变量维数；hidden：编码器、解码器每个隐藏层的宽度；D：像素个数（64） 【DGM 联系】编码器就是第 5 讲 §6 的变分分布 q(z; φ)，只是 φ 不再每张图一组，而是由网络从 x 算出来—— "摊销推断"（原课 slide 17–18）。解码器 p_θ(x | z) 就是第 5 讲 §2.4 的"无穷多个成分的混合"里的成分。 """ def __init__(self, d_z=2, hidden=256, D=64): super().__init__() self.d_z, self.D = d_z, D self.enc = nn.Sequential(nn.Linear(D, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 2 * d_z)) self.dec = nn.Sequential(nn.Linear(d_z, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, D)) def encode(self, x): """q_φ(z | x) 的参数：返回 (μ, log σ²)，形状都是 (N, d_z)。""" return self.enc(x).chunk(2, dim=-1) def decode_logits(self, z): """p_θ(x | z) 的 64 个 logit；z 可以有任意多的前导维度，例如 (k, N, d_z) → (k, N, 64)。""" return self.dec(z) def log_px_given_z(self, x, z): """log p_θ(x | z) = Σ_i log Bernoulli(x_i; σ(logit_i))；x 会自动广播到 z 的前导维度。""" logits = self.decode_logits(z) return -F.binary_cross_entropy_with_logits(logits, x.expand_as(logits), reduction="none").sum(-1) def elbo(self, x, generator, k=1, beta=1.0): """ 每张图的 ELBO = E_q[log p(x | z)] − β · KL(q(z | x) ‖ p(z))，用 k 个重参数化样本估计第一项。 【参数】x：(N, 64)；generator：抽 ε 用；k：每张图几个样本；beta：KL 前面的系数（β = 1 才是真正的 ELBO） 【返回】字典 {"elbo": (N,), "rec": (N,), "kl": (N,)}（单位 nat）。对 θ、φ 都可以反向传播。 【思路】重参数化：z = μ + σ ⊙ ε，ε ~ N(0, I)。随机性全在 ε 里，μ、σ 是 φ 的确定性函数，梯度可以穿过采样这一步。 KL(N(μ, σ²) ‖ N(0, 1)) = ½ (μ² + σ² − 1 − log σ²)，每一维求和——闭式解，没有采样噪声。 """ mu, logvar = self.encode(x) eps = torch.randn((k,) + mu.shape, generator=generator, dtype=mu.dtype) z = mu + (0.5 * logvar).exp() * eps # (k, N, d_z) rec = self.log_px_given_z(x, z).mean(0) # (N,) kl = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar).sum(-1) return {"elbo": rec - beta * kl, "rec": rec, "kl": kl} def log_weights(self, x, k, generator): """ k 个重要性权重的对数 log w_j = log p(x | z_j) + log p(z_j) − log q(z_j | x)，z_j ~ q(z | x)。 【返回】(k, N) 张量 """ mu, logvar = self.encode(x) eps = torch.randn((k,) + mu.shape, generator=generator, dtype=mu.dtype) z = mu + (0.5 * logvar).exp() * eps return self.log_px_given_z(x, z) + std_normal_log_prob(z) - gaussian_log_prob(z, mu, logvar) def iwae(self, x, k, generator): """ IWAE 下界 L_k = log (1/k) Σ_j w_j（每张图一个值，nat）。k = 1 时它的期望就是 ELBO（第 5 讲 §4.3）。 【边界】k 很大时占用内存 ∝ k × N × 64；大 k 请用 log_px（分块计算）。 """ lw = self.log_weights(x, k, generator) return torch.logsumexp(lw, dim=0) - math.log(k) @torch.no_grad() def log_px(self, x, k, generator, chunk=250, batch=500): """ 用 k 个重要性样本（提议 = 编码器 q）估计 log p_θ(x)：log (1/k) Σ_j w_j。分块计算，k 可以很大（如 5000）。 【返回】(N,) 张量（nat）。它是 log p_θ(x) 的**下界的估计**（第 5 讲 §4.3），k → ∞ 时收敛到 log p_θ(x)。 """ out = [] for i in range(0, x.shape[0], batch): xb = x[i:i + batch] acc = None for s in range(0, k, chunk): lw = self.log_weights(xb, min(chunk, k - s), generator) m = torch.logsumexp(lw.double(), dim=0) acc = m if acc is None else torch.logaddexp(acc, m) out.append(acc - math.log(k)) return torch.cat(out) @torch.no_grad() def sample(self, n, generator): """祖先采样：z ~ N(0, I)，x_i ~ Bernoulli(σ(logit_i))。返回 (样本 (n, 64), 像素概率 (n, 64))。""" z = torch.randn(n, self.d_z, generator=generator) p = torch.sigmoid(self.decode_logits(z)) return (torch.rand(p.shape, generator=generator) < p).float(), p def n_params(self): return sum(p.numel() for p in self.parameters()) @torch.no_grad() def exact_log_px_2d(model, X, lim=5.0, n=400, batch=250): """ d_z = 2 时，在 z ∈ [−lim, lim]² 的 n × n 网格上数值积分，得到"精确的" log p_θ(x) = log ∫ p(z) p_θ(x | z) dz。 【思路】∫ f(z) dz ≈ Σ_格点 f(z_g) × 格子面积；在对数空间里算：logsumexp_g [log p(z_g) + log p(x | z_g)] + log(面积)。 先把全部 n² 个格点解码一次（得到 n² × 64 的 logit），再对每批图做一次矩阵乘法算出全部 (图, 格点) 的 log p(x | z)。 【返回】(N,) float64 张量（nat） 【边界】只适用于 d_z = 2。后验比格子还窄时会不准——用 n 加倍检查结果是否变化（自检里做了）。 先验在 |z| > 5 处的质量约 6e−7，截断的影响可以忽略。 """ if model.d_z != 2: raise ValueError("exact_log_px_2d 只适用于 d_z = 2") t = torch.linspace(-lim, lim, n, dtype=torch.float64) Z = torch.cartesian_prod(t, t) # (n², 2) log_area = 2 * math.log((t[1] - t[0]).item()) logits = model.decode_logits(Z.float()).double() # (n², 64) log_s = F.logsigmoid(logits) # log σ(l) log_1ms = F.logsigmoid(-logits) # log(1 − σ(l)) log_prior = std_normal_log_prob(Z) out = [] for i in range(0, X.shape[0], batch): xb = X[i:i + batch].double() ll = xb @ log_s.T + (1 - xb) @ log_1ms.T # (b, n²)：log p(x | z_g) out.append(torch.logsumexp(ll + log_prior, dim=1) + log_area) return torch.cat(out) def train_vae(model, train, val, steps=15000, batch=128, lr=2e-3, seed=0, objective="elbo", k=1, beta=1.0, warmup=0, log_every=1000, verbose=False): """ 训练 VAE：小批量 Adam（余弦学习率），每 log_every 步在验证集上记录 ELBO（k = 1，β = 1）。 【参数】 objective："elbo"（最大化 ELBO）或 "iwae"（最大化 IWAE 下界 L_k） k：每张图的样本数（ELBO 时是重构项的样本数；IWAE 时是 L_k 的 k） beta：ELBO 里 KL 的系数（β-VAE）；warmup：前 warmup 步把 β 从 0 线性增加到 beta（"KL 退火"） 【返回】字典：history = [(步数, 验证集 −ELBO bits, 验证集 KL bits)]、seconds 【边界】β ≠ 1 时训练的不再是 log p(x) 的下界；评估时仍用 β = 1 的 ELBO。 """ g = torch.Generator().manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 0.5 * (1 + math.cos(math.pi * s / steps))) history, t0 = [], time.perf_counter() for step in range(1, steps + 1): x = train[torch.randint(0, train.shape[0], (batch,), generator=g)] b = beta * min(1.0, step / warmup) if warmup > 0 else beta if objective == "elbo": loss = -model.elbo(x, g, k=k, beta=b)["elbo"].mean() else: loss = -model.iwae(x, k, g).mean() opt.zero_grad() loss.backward() opt.step() sched.step() if step % log_every == 0 or step == steps: with torch.no_grad(): r = model.elbo(val, torch.Generator().manual_seed(99), k=1) history.append((step, -r["elbo"].mean().item() / LN2, r["kl"].mean().item() / LN2)) if verbose: print(f" 步 {step:5d}：验证集 −ELBO = {history[-1][1]:.3f} bit，KL = {history[-1][2]:.3f} bit") return {"history": history, "seconds": time.perf_counter() - t0} # 缓存的模型：名字 → (构造参数, 训练参数) TRAIN = {"steps": 15000, "lr": 2e-3, "batch": 128} # 每个约 10 秒（单线程 CPU） CONFIGS = {f"vae_z{d}": ({"d_z": d}, TRAIN) for d in (1, 2, 4, 8, 16)} CONFIGS["vae_z2_iwae10"] = ({"d_z": 2}, {**TRAIN, "objective": "iwae", "k": 10}) # 04 号脚本：IWAE 训练（约 25 秒） for _b in (0.25, 0.5, 2.0, 4.0): # 04 号脚本：β-VAE CONFIGS[f"vae_z8_beta{_b:g}"] = ({"d_z": 8}, {**TRAIN, "beta": _b}) def get_vae(name, data, verbose=False, quick=False): """ 取一个训练好的 VAE：checkpoints/name.pt 存在就直接读，否则按 CONFIGS[name] 训练并保存。 【参数】data：ar_models.pixel_data() 的返回值；quick=True：只训练 200 步、不保存（冒烟测试） 【返回】(model, info)：info 里有 history、seconds（训练用时）、cached 【边界】改了 VAE 的结构或 CONFIGS 以后，要删掉 checkpoints/ 里对应的文件。 """ kw_model, kw_train = CONFIGS[name] path = os.path.join(CKPT_DIR, f"{name}.pt") torch.manual_seed(1234) # 紧挨着建模型之前设种子：初始权重可复现 model = VAE(**kw_model) if os.path.exists(path) and not quick: state = torch.load(path, weights_only=False) model.load_state_dict(state["model"]) return model, {"history": state["history"], "seconds": state["seconds"], "cached": True} train, val = data[0], data[1] kw = dict(kw_train) if quick: kw["steps"] = 200 info = train_vae(model, train, val, verbose=verbose, **kw) if not quick: os.makedirs(CKPT_DIR, exist_ok=True) torch.save({"model": model.state_dict(), "history": info["history"], "seconds": info["seconds"]}, path) info["cached"] = False return model, info def _selftest(): torch.set_num_threads(1) g = torch.Generator().manual_seed(0) mu = torch.randn(5, 3, generator=g) logvar = torch.randn(5, 3, generator=g) z = mu + (0.5 * logvar).exp() * torch.randn(200000, 5, 3, generator=g) kl_mc = (gaussian_log_prob(z, mu, logvar) - std_normal_log_prob(z)).mean(0) kl_cf = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar).sum(-1) print(f"KL(q ‖ p) 闭式解 vs 20 万个样本的蒙特卡洛：最大差 {(kl_mc - kl_cf).abs().max().item():.4f}（应在 0.01 量级）") A = load_l3() data = A.pixel_data() model, info = get_vae("vae_z2", data) print(f"vae_z2：{'读取缓存' if info['cached'] else '训练完成'}（训练用时 {info['seconds']:.1f} 秒），参数 {model.n_params():,}") X = data[2][:200] with torch.no_grad(): exact = exact_log_px_2d(model, X).mean().item() exact2 = exact_log_px_2d(model, X, n=800).mean().item() rows = [(k, model.log_px(X, k, torch.Generator().manual_seed(1)).mean().item()) for k in [1, 10, 100, 1000]] print(f"网格积分 log p(x)：n = 400 → {exact / LN2:.4f} bit，n = 800 → {exact2 / LN2:.4f} bit") print("IWAE 估计随 k 上升、逼近网格积分：" + "，".join(f"k={k}: {v / LN2:.3f}" for k, v in rows)) if __name__ == "__main__": _selftest()