""" gan_lab.py —— 第 6 讲的 GAN 工具箱：二维玩具 GAN、DCGAN、条件 GAN 和训练循环 ===================================================================== 【这个文件要回答的问题】 GAN 由两个确定性的神经网络组成（slides）： 生成器 Generator G_θ(z)：输入一个随机噪声向量 z ~ N(0, I)，输出一个样本（一个点 / 一张图）； 判别器 Discriminator D_φ(x)：输入一个样本，输出"它是真的"的概率（真 = 1，假 = 0）。 训练是一个两人的极小极大博弈：D 想把真假分清，G 想骗过 D。这个文件把两个网络和"交替训练"写出来。 【本文件提供】 二维玩具： ring_data / RING 八个小高斯排成一圈（前置讲 P2 的 ring_of_gaussians），GAN 模式崩塌的经典测试 MLPGenerator / MLPDiscriminator mode_stats 样本落在了几个模式上（模式覆盖）、多少样本是"高质量"的 ascii_scatter 在终端里画二维散点的密度 图像： DCGenerator DCGAN 式生成器：全连接 → 7×7 → 转置卷积 → 14×14 → 转置卷积 → 28×28 DCDiscriminator 卷积判别器：28×28 → 14×14 → 7×7 → 一个 logit （两者都可以传 n_classes 变成条件 GAN：把类别嵌入接到生成器的输入、把类别"平面"接到判别器的输入） 训练： d_loss / g_loss 判别器和生成器的损失（g_loss 支持 slides 的"极小极大"版本和常用的"非饱和"版本） train_gan 通用的交替训练循环（二维和图像共用） 【约定】 - 判别器输出的是 logit a（还没过 sigmoid），D(x) = σ(a)。用 logit 计算损失在数值上更稳定： log σ(a) = −softplus(−a)，log(1 − σ(a)) = −softplus(a)，不会出现 log(0)。 - 所有随机性来自显式的 torch.Generator 和 torch.manual_seed：同一台机器上重复运行逐位相同。 【GenAI 联系】 GAN 是 2014–2021 年图像生成的主角（DCGAN、StyleGAN、BigGAN），后来被扩散模型取代（第 7–8 讲）； 但"判别器当作可学习的损失函数"这个想法仍在：Stable Diffusion 的 VAE 解码器、SDXL-Turbo 的一步生成、 很多语音合成模型都用对抗损失（第 12、14、24 讲）。 """ import math import os import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) _P2 = os.path.join(HERE, "..", "prereq2") if _P2 not in sys.path: sys.path.insert(0, _P2) from prob_lab import ring_of_gaussians # noqa: E402 RING = ring_of_gaussians(n_modes=8, radius=2.0, std=0.05) # 前置讲 P2 的同一个分布（GMM2D 对象） # ===================================================================== # 1. 二维玩具：数据、网络、评估 # ===================================================================== def ring_data(n, rng): """从八高斯圆环采 n 个点。【参数】rng：np.random.default_rng(种子)。【返回】(n, 2) 的 float32 张量。""" return torch.from_numpy(RING.sample(n, rng)).float() class MLPGenerator(nn.Module): """z (z_dim) → [Linear → ReLU] × n_layers → Linear → 输出 (out_dim)。最后一层不加激活：输出可以是任意实数。""" def __init__(self, z_dim=8, hidden=128, n_layers=3, out_dim=2): super().__init__() layers, d = [], z_dim for _ in range(n_layers): layers += [nn.Linear(d, hidden), nn.ReLU()] d = hidden self.net = nn.Sequential(*layers, nn.Linear(d, out_dim)) self.z_dim = z_dim def forward(self, z): return self.net(z) class MLPDiscriminator(nn.Module): """x (in_dim) → [Linear → LeakyReLU(0.2)] × n_layers → Linear → 一个 logit。""" def __init__(self, in_dim=2, hidden=128, n_layers=3): super().__init__() layers, d = [], in_dim for _ in range(n_layers): layers += [nn.Linear(d, hidden), nn.LeakyReLU(0.2)] d = hidden self.net = nn.Sequential(*layers, nn.Linear(d, 1)) def forward(self, x): return self.net(x).squeeze(-1) def mode_stats(x, gmm=RING, n_std=3.0): """ 评估二维样本对八个模式的覆盖。 【思路】每个点归到最近的模式中心；离中心不到 n_std 个标准差（默认 3σ = 0.15）的点算"高质量"。 一个模式被"覆盖"：落在它附近的高质量样本至少占总样本的 1/32（均匀时每个模式应占 1/8，这里只要求四分之一）。 【参数】x：(n, 2) 张量或数组；gmm：GMM2D 对象 【返回】字典 {covered: 覆盖的模式数, hq: 高质量样本比例, counts: 每个模式的高质量样本数（列表）} """ x = np.asarray(x, dtype=np.float64) k = gmm.nearest_mode(x) dist = np.linalg.norm(x - gmm.means[k], axis=1) hq = dist < n_std * gmm.std counts = np.bincount(k[hq], minlength=len(gmm.means)) return {"covered": int((counts >= len(x) / 32).sum()), "hq": float(hq.mean()), "counts": counts.tolist()} def ascii_scatter(x, lim=2.6, width=41, height=21, chars=" .:+*#@"): """ 把二维点画成字符密度图：把 [−lim, lim]² 切成 height × width 个格子，格子里点越多字符越"重"。 【返回】字符串列表（每行一个），上方是 y 的正方向。 【边界】范围以外的点不画（但会算进 mode_stats）。 """ x = np.asarray(x, dtype=np.float64) ix = np.floor((x[:, 0] + lim) / (2 * lim) * width).astype(int) iy = np.floor((lim - x[:, 1]) / (2 * lim) * height).astype(int) ok = (ix >= 0) & (ix < width) & (iy >= 0) & (iy < height) grid = np.zeros((height, width), dtype=int) np.add.at(grid, (iy[ok], ix[ok]), 1) # 按对数刻度分级：1 个点就有显示，几百个点才到最重的字符 level = np.where(grid > 0, np.minimum(len(chars) - 1, 1 + np.log2(np.maximum(grid, 1)).astype(int) // 2), 0) return ["".join(chars[v] for v in row) for row in level] # ===================================================================== # 2. 图像：DCGAN 式的生成器和判别器 # ===================================================================== class DCGenerator(nn.Module): """ DCGAN 式生成器（slides："一个倒过来的 CNN，用分数步长卷积把图像一层层变大"）。 z (z_dim) [+ 类别嵌入] → Linear → (2c, 7, 7) → BN → ReLU → ConvTranspose2d(2c → c, 核 4, 步长 2, 填充 1) → (c, 14, 14) → BN → ReLU → ConvTranspose2d(c → 1, 核 4, 步长 2, 填充 1) → (1, 28, 28) → Sigmoid（像素值在 (0, 1)） 【转置卷积 = 分数步长卷积】步长 2 的转置卷积相当于先在输入像素之间插 0（"步长 1/2"），再做普通卷积， 所以尺寸翻倍：输出 = (输入 − 1) × 步长 − 2 × 填充 + 核 = (7 − 1) × 2 − 2 + 4 = 14。03 号脚本第 0 部分验证了这一点。 【参数】z_dim：噪声维数；c：通道数；n_classes：> 0 时变成条件生成器（类别嵌入维数 = 8，接在 z 后面） """ def __init__(self, z_dim=32, c=32, n_classes=0): super().__init__() self.z_dim, self.n_classes = z_dim, n_classes self.embed = nn.Embedding(n_classes, 8) if n_classes else None d_in = z_dim + (8 if n_classes else 0) self.fc = nn.Sequential(nn.Linear(d_in, 2 * c * 7 * 7), nn.BatchNorm1d(2 * c * 7 * 7), nn.ReLU()) self.up = nn.Sequential( nn.ConvTranspose2d(2 * c, c, 4, 2, 1), nn.BatchNorm2d(c), nn.ReLU(), nn.ConvTranspose2d(c, 1, 4, 2, 1), nn.Sigmoid()) self.c = c def forward(self, z, y=None): if self.n_classes: z = torch.cat([z, self.embed(y)], dim=1) return self.up(self.fc(z).view(-1, 2 * self.c, 7, 7)) class DCDiscriminator(nn.Module): """ 卷积判别器：(1, 28, 28) → Conv(核 4, 步长 2) → (c, 14, 14) → LeakyReLU → Conv(核 4, 步长 2) → (2c, 7, 7) → LeakyReLU → 拉直 → Linear → 一个 logit。 【为什么用步长 2 的卷积而不是池化】DCGAN 论文的建议：让网络自己学怎样下采样。 【为什么判别器不用 BatchNorm】BN 让一张图的输出依赖同一批的其他图；判别器一批里全是真图或全是假图时， 这会泄漏"这一批是真是假"的信息。这里干脆不用（另一种常见做法是真假混在一批里）。 【条件判别器】n_classes > 0 时，把类别的 one-hot 展开成 n_classes 张常数"平面"，和图像拼在通道维上 （slides："把标签嵌入接到生成器和判别器的输入上"）。 """ def __init__(self, c=16, n_classes=0): super().__init__() self.n_classes = n_classes self.net = nn.Sequential( nn.Conv2d(1 + n_classes, c, 4, 2, 1), nn.LeakyReLU(0.2), nn.Conv2d(c, 2 * c, 4, 2, 1), nn.LeakyReLU(0.2), nn.Flatten(), nn.Linear(2 * c * 7 * 7, 1)) def forward(self, x, y=None): if self.n_classes: planes = F.one_hot(y, self.n_classes).float()[:, :, None, None].expand(-1, -1, *x.shape[2:]) x = torch.cat([x, planes], dim=1) return self.net(x).squeeze(-1) # ===================================================================== # 3. 损失与训练循环 # ===================================================================== def d_loss(logit_real, logit_fake): """ 判别器的损失 = 二分类交叉熵（真 = 1，假 = 0），也就是 slides 里 max_φ log D(x) + log(1 − D(G(z))) 取负号： −log σ(a_real) − log(1 − σ(a_fake)) = softplus(−a_real) + softplus(a_fake) 【返回】一批的平均损失（标量）。完全分不清时（D 处处输出 0.5）等于 2 ln 2 ≈ 1.386。 """ return F.softplus(-logit_real).mean() + F.softplus(logit_fake).mean() def g_loss(logit_fake, kind="ns"): """ 生成器的损失。 kind="minimax"：slides 的原始目标 min_θ log(1 − D(G(z))) = −softplus(a)。 问题：训练初期 D 很容易认出假图，D(G(z)) ≈ 0，此时它对 a 的梯度是 −σ(a) ≈ 0——生成器几乎学不到东西（"饱和"）。 kind="ns"（non-saturating，非饱和）：改成 max_θ log D(G(z))，即最小化 −log σ(a) = softplus(−a)。 它和原目标有同样的不动点，但在 D(G(z)) ≈ 0 时梯度是 −(1 − σ(a)) ≈ −1，信号最强。 Goodfellow et al. (2014) 论文里就建议这样做，实践中几乎都用它。01 号脚本第 2 部分有数值对比。 """ if kind == "minimax": return -F.softplus(logit_fake).mean() if kind == "ns": return F.softplus(-logit_fake).mean() raise ValueError(kind) def train_gan(G, D, sample_real, steps, batch=64, lr_g=2e-4, lr_d=2e-4, betas=(0.5, 0.999), kind="ns", d_steps=1, seed=0, callback=None, every=0, labels=None): """ 交替训练（slides：固定 G、对 D 反向传播；固定 D、对 G 反向传播；不精确求解极小极大，每次只走一步小批量 SGD）。 【参数】 G, D 生成器和判别器 sample_real 函数 sample_real(batch, gen) → 真样本；条件 GAN 时返回 (样本, 标签) steps 生成器的更新次数 batch 每步的样本数（真假各 batch 个） lr_g, lr_d, betas 两个 Adam 优化器的参数（β₁ = 0.5 是 DCGAN 论文的建议：动量太大时博弈容易振荡） kind 生成器损失："ns" 或 "minimax"（见 g_loss） d_steps 每更新一次生成器，先更新几次判别器（Goodfellow 2014 的算法 1 里的 k） callback 每 every 步调用一次 callback(step, G, D, stats)，用来记录训练过程 labels 条件 GAN 的类别数（None 表示无条件） 【返回】最后 100 步的平均 (d_loss, g_loss) 【边界】G 和 D 在函数里被切换到 train 模式，返回前切回 eval 模式（BatchNorm 在 eval 模式下用累计的均值方差）。 """ gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) opt_g = torch.optim.Adam(G.parameters(), lr=lr_g, betas=betas) opt_d = torch.optim.Adam(D.parameters(), lr=lr_d, betas=betas) G.train() D.train() hist = [] for step in range(1, steps + 1): # ---- 1. 判别器：固定 G（生成的假样本 detach，梯度不传回 G），最小化二分类交叉熵 ---- for _ in range(d_steps): if labels: xr, yr = sample_real(batch, gen) yf = torch.randint(labels, (batch,), generator=gen) fake = G(torch.randn(batch, G.z_dim, generator=gen), yf).detach() ld = d_loss(D(xr, yr), D(fake, yf)) else: xr = sample_real(batch, gen) fake = G(torch.randn(batch, G.z_dim, generator=gen)).detach() ld = d_loss(D(xr), D(fake)) opt_d.zero_grad() ld.backward() opt_d.step() # ---- 2. 生成器：固定 D（只更新 G 的参数；梯度穿过 D 流回 G），让 D 把假样本判成真 ---- z = torch.randn(batch, G.z_dim, generator=gen) if labels: yf = torch.randint(labels, (batch,), generator=gen) lg = g_loss(D(G(z, yf), yf), kind) else: lg = g_loss(D(G(z)), kind) opt_g.zero_grad() lg.backward() opt_g.step() # 只走 G 的优化器：D 的参数虽然也算出了梯度，但下一次 opt_d.zero_grad() 会清掉 hist.append((ld.item(), lg.item())) if callback is not None and every and step % every == 0: G.eval() callback(step, G, D, hist) G.train() G.eval() D.eval() tail = hist[-100:] return sum(h[0] for h in tail) / len(tail), sum(h[1] for h in tail) / len(tail) @torch.no_grad() def generate(G, n, seed, y=None, batch=500): """用固定种子从 G 采 n 个样本（分批，避免一次占太多内存）。y：条件 GAN 的类别（(n,) 张量）。""" gen = torch.Generator().manual_seed(seed) G.eval() out = [] for i in range(0, n, batch): z = torch.randn(min(batch, n - i), G.z_dim, generator=gen) out.append(G(z) if y is None else G(z, y[i:i + batch])) return torch.cat(out) def count_parameters(model): return sum(p.numel() for p in model.parameters()) def jsd(p, q): """ 两个离散分布（同一个网格上的概率向量）之间的 Jensen–Shannon 散度（nat）： JSD(p ‖ q) = ½ KL(p ‖ m) + ½ KL(q ‖ m)，m = (p + q)/2。 【性质】对称；0 ≤ JSD ≤ ln 2；两个分布不重叠时恰好等于 ln 2（不管它们相距多远）。 """ p, q = np.asarray(p, dtype=np.float64), np.asarray(q, dtype=np.float64) m = (p + q) / 2 def kl(a, b): nz = (a > 0) & (b > 0) # 数学上 a > 0 时 m ≥ a/2 > 0；b == 0 只可能是极小的数下溢，这些项可以忽略 return float(np.sum(a[nz] * np.log(a[nz] / b[nz]))) return 0.5 * kl(p, m) + 0.5 * kl(q, m) LOG4 = math.log(4.0)