从原理走到生成：扩散模型笔记

扩散模型 图像生成 PyTorch 采样策略
作者: 醉月思 发布日期: 2023-01-17 分类: 技术实践

前阵子扩散模型笔记连续改了几轮，趁还记得写成备忘。

前向过程：把图变成噪声

扩散模型的前向过程本质上是马尔可夫链：从真实数据分布 $x_0$ 开始，一步步加入高斯噪声，经过 $T$ 步后，$x_T$ 近似服从标准正态分布。

每一步的噪声量由 $\beta_t$ 控制，通常设计为从小到大递增的序列，让噪声"一点点"加进来。实际实现时，我们用重参数化技巧直接采样 $x_t$：

$$ x_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1 - \bar\alpha_t} \epsilon $$

其中 $\epsilon \sim \mathcal{N}(0, I)$ 是标准高斯噪声，$\bar\alpha_t = \prod_{i=1}^t (1 - \beta_i)$。

为什么这么写？因为这避免了 $T$ 步循环。给定 $x_0$ 和任意 $t$，我们可以直接得到 $x_t$，这是训练和采样时的关键优化。

import torch

def q_sample(x_0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)
    
    sqrt_alpha_bar = torch.sqrt(alpha_bar[t])
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bar[t])
    
    return sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise

这里 alpha_bar 是预计算好的累积乘积，t 是时间步索引。代码很直白，但调试时遇到过两个坑：

  1. 时间步越界：t 如果是 1000 而 alpha_bar 长度是 1000，会直接 IndexError。训练时 t 通常用 torch.randint 采样，记得上限要减 1。
  2. 形状不匹配：如果 x_0 是 (B, C, H, W)，t 是标量，这里没问题；但如果 t 是一个 (B,) 的张量，就需要加维度 t.view(-1, 1, 1, 1) 才能广播。这个坑在写 DDPM 原论文的实现时经常碰到。

逆向过程：从噪声学到去噪

前向过程是固定的，不需要训练。真正要学的是逆向过程：给定 $x_t$，预测噪声 $\epsilon_\theta(x_t, t)$，然后用这个预测算出 $x_{t-1}$。

训练目标很简单：给定 $x_0$ 和采样 $t$，生成噪声 $\epsilon$，用模型预测 $\epsilon_\theta(x_t, t)$，然后最小化两者的 MSE。

def p_losses(model, x_0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)
    
    x_t = q_sample(x_0, t, noise)
    predicted_noise = model(x_t, t)
    
    return torch.nn.functional.mse_loss(noise, predicted_noise)

这里的 model 是一个 UNet，接受图像和时间步输入，输出预测噪声。UNet 的结构不展开说，但实现时注意几点：

  1. 时间步嵌入：t 通常被转成 embedding 再注入到各个残差块，可以用 sinusoidal embedding 或可学习的 embedding。我试过两者， sinusoidal 在小模型上稳定，可学习在大的迭代次数下稍好一点。
  2. 注意力的使用：UNet 中间层可以加自注意力，帮助捕捉全局依赖，但内存开销明显。实践时从 16x16 分辨率开始加注意力，再低就扛不住了。
  3. 时间步如何传递：如果 model 接收的是标量时间步，内部要记得归一化到 [0, 1] 或者对应 embedding 的范围，否则数值太大会让后续层爆炸。

采样：去噪链的实战

训练完后，采样就是从 $x_T \sim \mathcal{N}(0, I)$ 开始，一步步"后退"到 $x_0$。每一步的均值和方差都可以推导出来：

$$ \mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}} \epsilon_\theta(x_t, t) \right) $$

方差通常用固定值 $\sigma_t^2 = \beta_t$，简化实现。

@torch.no_grad()
def p_sample(model, x_t, t):
    predicted_noise = model(x_t, t)
    alpha_t = alpha[t]
    beta_t = beta[t]
    alpha_bar_t = alpha_bar[t]
    
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bar_t)
    sqrt_recip_alpha = torch.sqrt(1 / alpha_t)
    
    pred_mean = sqrt_recip_alpha * (x_t - beta_t / sqrt_one_minus_alpha_bar * predicted_noise)
    
    if t == 0:
        return pred_mean
    else:
        noise = torch.randn_like(x_t)
        return pred_mean + torch.sqrt(beta_t) * noise

这里有个常见的坑：t=0 时不应再加噪声，否则生成图永远不会"干净"。实现时一定要检查时间步边界。

完整采样循环就是从 T-1 倒序到 0：

@torch.no_grad()
def p_sample_loop(model, shape):
    device = next(model.parameters()).device
    img = torch.randn(shape, device=device)
    
    for i in reversed(range(T)):
        img = p_sample(model, img, torch.full((shape[0],), i, device=device, dtype=torch.long))
    
    return img

调试时遇到两个问题：

  1. 采样太慢：如果 T=1000，每步都要前向推理一次，生成一张图要几分钟。可以尝试减少 T 到 250 或 100，损失一点质量换取速度。我试过 250 步，MNIST 上肉眼差别不大。
  2. 数值不稳定：有时中间几步的 pred_mean 会突然爆炸，最后生成全白或全黑。这通常是因为 beta_t 序列设计得不好，或者时间步 embedding 有问题。改用 cosine schedule 的 beta_t 后，问题缓解不少。

真实环境与踩坑记录

我用 PyTorch 在 MNIST 和 CIFAR-10 上跑过基础的 DDPM，主要配置：

  • 框架：PyTorch 2.1.0，CUDA 12.1
  • 模型：UNet-base，参数量约 30M（CIFAR-10）和 10M（MNIST）
  • 训练轮数：MNIST 100 epochs，CIFAR-10 300 epochs
  • T=1000，线性 beta_t 从 1e-4 到 2e-2
  • 优化器：AdamW，lr=2e-4，weight_decay=1e-4
  • Batch size：256（MNIST），64（CIFAR-10，显存吃紧）

踩坑 1：显存不够用

训练 CIFAR-10 时，64 的 batch size 在 24GB 显存上只能勉强塞下。加注意力后直接 OOM。解决方案：

  1. 用梯度检查点（torch.utils.checkpoint）换时间换空间。
  2. 减少 UNet 的通道数，从 128 降到 64。
  3. 用混合精度训练（torch.cuda.amp），显存占用能降 30% 左右。

梯度检查点会让训练慢 20% 左右，但在资源有限时是必要的折衷。

踩坑 2：采样收敛慢

训练完成后，前几十步采样出来的图基本是噪声，要到倒数 100 步才开始"像样"。这说明学习到的去噪能力在晚期才发挥作用。尝试了几个方案：

  1. 重训练：换用 cosine schedule 的 beta_t，早期噪声加得更少，让模型多接触"清晰"的中间态。效果稍有改善，但不明显。
  2. 改采样器：从 DDPM 换到 DDIM，采样步数可以减到 50，但质量下降明显。最后还是老老实实跑 1000 步。
  3. 早停条件：在采样循环中检测方差，如果预测噪声很小就提前退出。有时能省 30% 时间，但偶尔会卡在"半成品"。

实践下来，如果只追求能跑通，线性 schedule 加完整 1000 步是最稳定的。要追求质量，还得从模型本身入手。

踩坑 3：图像质量差

初始生成的 MNIST 数字还勉强能认，但 CIFAR-10 的图基本是"带颜色的噪声"。排查后发现几个原因：

  1. 训练时间不够：300 epochs 后，loss 还在缓慢下降。跑满 500 epochs 后，质量有肉眼可见提升。
  2. 模型容量不足：30M 参数的 UNet 对 CIFAR-10 还是偏小。换到 100M 后，生成的图开始有"轮廓"，但训练时间翻倍。
  3. 评估指标误导：FID 计算需要大量样本，小 batch 下算不准。不要因为 FID 突然变好或变坏就急着调参。

完整可运行示例

下面是一个简化的、能在 MNIST 上跑通的最小实现。代码去掉了很多工程细节，但保留了核心逻辑：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from tqdm import tqdm

# 配置
T = 1000
beta_start = 0.0001
beta_end = 0.02
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# beta_t 线性递增
beta = torch.linspace(beta_start, beta_end, T, device=device)
alpha = 1 - beta
alpha_bar = torch.cumprod(alpha, dim=0)

# 重参数化采样
def q_sample(x_0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)
    sqrt_alpha_bar = torch.sqrt(alpha_bar[t])
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bar[t])
    return sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise

# 损失函数
def p_losses(model, x_0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)
    x_t = q_sample(x_0, t, noise)
    predicted_noise = model(x_t, t)
    return torch.nn.functional.mse_loss(noise, predicted_noise)

# 简化的 UNet (只保留主干)
class SimpleUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.time_embed = nn.Sequential(
            nn.Linear(1, 128),
            nn.SiLU(),
            nn.Linear(128, 128)
        )
        self.down = nn.Sequential(
            nn.Conv2d(1, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU()
        )
        self.up = nn.Sequential(
            nn.Conv2d(64, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU(),
            nn.Conv2d(64, 1, 3, padding=1)
        )
    
    def forward(self, x, t):
        t_emb = self.time_embed(t.float().unsqueeze(1) / T)
        t_emb = t_emb.view(-1, 128, 1, 1)
        x = self.down(x)
        x = x + t_emb  # 简单的时间步注入
        x = self.up(x)
        return x

# 采样
@torch.no_grad()
def p_sample(model, x_t, t):
    predicted_noise = model(x_t, t)
    alpha_t = alpha[t]
    beta_t = beta[t]
    alpha_bar_t = alpha_bar[t]
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bar_t)
    sqrt_recip_alpha = torch.sqrt(1 / alpha_t)
    pred_mean = sqrt_recip_alpha * (x_t - beta_t / sqrt_one_minus_alpha_bar * predicted_noise)
    if t == 0:
        return pred_mean
    else:
        noise = torch.randn_like(x_t)
        return pred_mean + torch.sqrt(beta_t) * noise

@torch.no_grad()
def p_sample_loop(model, shape):
    img = torch.randn(shape, device=device)
    for i in reversed(range(T)):
        img = p_sample(model, img, torch.full((shape[0],), i, device=device, dtype=torch.long))
    return img

# 训练
def train():
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True)
    
    model = SimpleUNet().to(device)
    optimizer = optim.AdamW(model.parameters(), lr=2e-4)
    
    model.train()
    for epoch in range(100):
        for x_0, _ in tqdm(dataloader, desc=f"Epoch {epoch}"):
            x_0 = x_0.to(device)
            t = torch.randint(0, T, (x_0.shape[0],), device=device).long()
            loss = p_losses(model, x_0, t)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    torch.save(model.state_dict(), "ddpm_mnist.pt")

if __name__ == "__main__":
    train()

这段代码在 MNIST 上训练 100 个 epoch 后，p_sample_loop 能生成出可识别的数字，虽然细节模糊，但足以说明扩散模型的基本流程。

小结与余味

扩散模型看起来复杂，核心思想却很朴素：先把图"毁"成噪声，再学着"修"回去。前向过程固定，逆向可学，训练目标直白，采样过程也不玄乎。

真正麻烦的是工程细节：怎么设计 beta_t 序列、怎么平衡模型容量与显存、怎么加速采样而不牺牲质量。这些书本上写得少，实践里撞得多。

写这篇文章时，我翻了不少论文和实现，最深的感受是：理论再漂亮，也得住代码里跑一遍才知道哪里会断。扩散模型的热度会过去，但这种"从直觉到实现"的学习方式，应该是程序员长期的本钱。

下次再看到那些让人眼花的公式，不妨先问一句：如果我直接写代码，最简单的那一行应该是什么？

版权声明: 本文首发于 指尖魔法屋-从原理走到生成：扩散模型笔记（https://blog.thinkmoon.cn/post/170-diffusion-models-practice-to-generation/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章

AI_Top-k采样踩坑记录
调对话生成模型时，我碰到一个很烦的现象：同一个问题问三遍，措辞几乎一模一样，连举例子的顺序都不变。
📁 技术实践
📅 2025-12-27