1. 从噪声到图像DDPM的核心思想与复现价值如果你对生成式AI感兴趣那么“扩散模型”这个词你一定不陌生。从DALL-E 2到Stable Diffusion这些能根据文字描述生成逼真图像的“魔法”背后核心引擎就是扩散模型。而DDPM即去噪扩散概率模型可以说是现代扩散模型的奠基之作。它用一套极其优雅的数学框架教会了神经网络如何从纯粹的随机噪声中一步步“雕刻”出我们想要的图像。网上关于DDPM原理的解读很多公式推导也足够深入但当你真正想动手实现一个时往往会发现理论与代码之间存在一道鸿沟。那些论文里的“前向过程”、“反向过程”、“变分下界”如何用PyTorch的Tensor操作表达出来时间步长t该如何嵌入到网络里那个关键的UNet结构到底长什么样每一层为什么要这么设计扩散调度中的β_t序列又该如何影响训练和采样这篇文章我将带你从零开始用PyTorch完整复现一个DDPM。我们不只停留在调用高级API而是要亲手搭建时间嵌入模块、构建UNet网络、实现扩散调度器并理解其中每一个超参数和设计选择的用意。这个过程就像组装一台精密的钟表你会看清每一个齿轮模块是如何咬合最终让整个系统运转起来。无论你是想深入理解扩散模型的内在机理还是为后续研究更复杂的模型如带条件的扩散模型打下坚实基础这次“造轮子”的经历都将让你获益匪浅。2. 扩散模型的“剧本”前向与反向过程详解在搭建任何模型之前我们必须彻底理解它要解决什么问题以及解决问题的“剧本”是什么。DDPM的剧本分为截然相反的两幕破坏与重建。2.1 前向过程给图像逐步加噪第一幕是前向过程或者叫扩散过程。这个过程是确定的、无需学习的。想象我们有一张清晰的原始图片x_0。我们准备了一个包含T个时间步的“噪声添加计划”。在每一个时间步t我们向图片中添加一点点高斯噪声。具体来说我们从x_{t-1}得到x_t的公式是x_t sqrt(1 - β_t) * x_{t-1} sqrt(β_t) * ε_t其中ε_t是从标准正态分布中采样的噪声β_t是一个预先定义好的、很小的正数例如从0.0001线性增长到0.02它决定了每一步添加的噪声量。sqrt(1 - β_t)和sqrt(β_t)是为了保证每一步的方差得到合理控制。这个过程有一个非常美妙的性质由于每一步都是线性高斯变换我们可以通过重参数化技巧直接从原始图像x_0计算出任意中间时刻t的加噪图像x_t而无需一步步迭代。公式如下x_t sqrt(ᾱ_t) * x_0 sqrt(1 - ᾱ_t) * ε这里α_t 1 - β_t而ᾱ_t Π_{s1}^{t} α_s即所有α_s的连乘。ε同样是标准高斯噪声。这个公式是代码实现的关键它意味着我们可以在训练时随机采样一个时间步t然后直接用这个公式构造出对应的加噪样本x_t效率极高。前向过程的本质它是一个将数据分布清晰的图片逐渐转变为简单先验分布标准高斯噪声的过程。当T足够大时x_T就几乎完全是一团纯粹的随机噪声了。2.2 反向过程让神经网络学会去噪第二幕是反向过程这才是模型需要学习的部分。我们的目标是逆转前向过程给定一个噪声图像x_t和时间步t我们想要预测出那一步所添加的噪声ε或者等价地预测出更清晰的上一时刻图像x_{t-1}。DDPM选择让神经网络ε_θ去预测噪声。为什么是预测噪声而不是直接预测x_{t-1}呢这背后有数学上的考量与变分下界的推导有关但从直觉上理解也很直接预测一个均值为0、方差为1的标准高斯噪声分量对于神经网络来说是一个更稳定、更容易学习的任务。因此反向过程的核心步骤是向神经网络输入加噪图像x_t和时间步t。神经网络输出对噪声ε的预测ε_θ(x_t, t)。利用预测的噪声我们可以根据以下公式估算出x_{t-1}x_{t-1} 1 / sqrt(α_t) * (x_t - (1 - α_t)/sqrt(1 - ᾱ_t) * ε_θ(x_t, t)) σ_t * z其中z是额外的随机噪声当t1时加入以增加生成样本的多样性σ_t是其方差。通过从tT到t1迭代执行上述步骤我们就能从一团噪声x_T ~ N(0, I)开始逐步“去噪”最终得到一张清晰的图像x_0。理解这个“剧本”的价值它决定了我们代码的整体框架。我们的模型UNet扮演的是“噪声预测器”的角色而训练目标就是最小化预测噪声ε_θ与真实添加噪声ε之间的差距。接下来我们就来搭建这位核心演员——UNet。3. 构建核心演员时间感知的UNet网络UNet是DDPM中用于噪声预测的骨干网络其结构像一个“U”形先下采样编码提取特征再上采样解码重建图像。但DDPM中的UNet有一个关键变体它必须能感知“时间”。因为去噪的动作强烈依赖于当前处于去噪过程的哪一步时间步t网络需要知道“我现在是在清理早期的大块噪声还是在做最后的精修”。3.1 时间步嵌入让网络拥有“时间感”时间步t是一个标量整数如250我们需要将其转化为网络能够利用的特征。直接输入全连接层是一种方法但DDPM论文采用了Transformer中类似的正弦位置编码思路并将其通过一个小的MLP多层感知机进行变换。import math import torch import torch.nn as nn class TimeEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim dim # 构建一个包含两个线性层和SiLU激活的小型MLP self.mlp nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), # SiLU (Swish)激活函数实践中比ReLU在扩散模型中效果更好 nn.Linear(dim * 4, dim) ) def forward(self, t): # t: (batch_size, 1) 整数时间步 half_dim self.dim // 2 # 生成频率因子和Transformer的位置编码类似 freqs torch.exp( -math.log(10000) * torch.arange(start0, endhalf_dim, dtypetorch.float32) / half_dim ).to(t.device) # 将t与频率因子相乘得到角度参数 args t.float() * freqs # 分别计算正弦和余弦部分然后拼接得到位置编码 embedding torch.cat([torch.sin(args), torch.cos(args)], dim-1) # 如果嵌入维度是奇数进行填充以匹配维度 if self.dim % 2: embedding torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim-1) # 通过MLP进行非线性变换得到最终的时间嵌入向量 return self.mlp(embedding)这个TimeEmbedding模块的输出是一个(batch_size, dim)的张量。接下来我们需要将这个时间信息注入到UNet的每一层中。3.2 UNet的层结构残差块与注意力机制一个标准的DDPM UNet由多个“下采样块”和“上采样块”组成。每个块的核心是一个“残差块”它负责主要的特征变换。此外在网络的深层特征图尺寸较小处会插入“注意力块”让模型能够建立图像不同区域之间的长程依赖这对于生成结构连贯的图像至关重要。残差块的设计它需要融合三种信息输入特征图x、时间嵌入t_emb以及在某些情况下对于条件生成的条件信息。典型的做法是对输入特征x进行GroupNorm和SiLU激活。将时间嵌入t_emb通过一个线性层投影到与特征图通道数匹配的维度然后加到特征上通常是在某个归一化层之后。使用一个卷积层进行特征变换。如果该块有残差连接即输入输出通道数相同则将变换后的特征与输入相加。注意力块的引入在特征图尺寸为16x16或8x8的层级我们会使用多头自注意力机制。由于此时空间维度已大大降低计算注意力权重的开销是可接受的。注意力机制能让模型在去噪时更好地协调图像不同部分的内容例如确保生成的猫有两只对称的眼睛。3.3 完整的UNet组装与下/上采样UNet的整体流程如下初始卷积将输入的(batch, 3, H, W)图像映射到初始通道数如128。编码器路径下采样由多个阶段组成。每个阶段包含若干个残差块每个块都注入当前时间步信息。可能的一个注意力块在深层。一个下采样层通常用步长为2的卷积或池化将特征图尺寸减半通道数翻倍。瓶颈层在最低分辨率下使用多个连续的残差块和注意力块进行深度处理。解码器路径上采样与编码器对称。每个阶段包含一个上采样层转置卷积或最近邻插值卷积将特征图尺寸翻倍通道数减半。若干个残差块。关键点解码器每个残差块的输入除了来自上一层的特征和当前时间嵌入还会拼接Concatenate来自编码器路径对应层级的特征。这是UNet的核心思想它让解码器在重建细节时能利用编码器提取的底层信息。最终卷积将通道数映射回3RGB并使用tanh或sigmoid激活将值约束到[-1, 1]或[0, 1]的像素范围。实操心得通道数与下采样倍数在设计UNet时一个常见的陷阱是下采样次数太多。对于64x64的图像下采样4次64-32-16-8-4是合理的。但如果下采样到1x1会丢失大量空间信息导致生成质量下降。同时初始通道数不宜过小如32通常128或256是一个好的起点以确保模型有足够的容量学习复杂的去噪映射。4. 导演的指令表扩散调度器的实现如果说UNet是演员那么扩散调度器就是导演手中的剧本和指令表。它严格定义了前向过程中每一步的噪声强度β_t并提供了反向采样时所需的全部系数。调度器的设计直接影响训练稳定性和生成质量。4.1 β_t 序列的生成线性与余弦调度最直接的策略是线性调度即让β_t从一个小值β_start线性增长到一个大值β_end。def linear_beta_schedule(timesteps, beta_start1e-4, beta_end0.02): return torch.linspace(beta_start, beta_end, timesteps)然而线性调度在过程的两端非常清晰和非常噪声变化可能过于剧烈。后续研究提出了改进的余弦调度它能让ᾱ_t即x_t中保留的原始信号系数随着t的变化遵循一个余弦函数的一半周期从而在过程开始和结束时变化更平缓。def cosine_beta_schedule(timesteps, s0.008): steps timesteps 1 x torch.linspace(0, timesteps, steps) alphas_cumprod torch.cos(((x / timesteps) s) / (1 s) * math.pi * 0.5) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] # 归一化使ᾱ_01 betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0, 0.999)在实际操作中余弦调度通常能带来更平滑的训练过程和更好的生成效果尤其是在步数较多时。4.2 关键系数的预计算一旦确定了β_t序列我们就可以预先计算出一系列在训练和采样中反复使用的系数并将其注册为模型的缓冲张量这能极大提升效率。# 假设 betas 是形状为 (T,) 的张量 alphas 1. - betas alphas_cumprod torch.cumprod(alphas, dim0) # ᾱ_t alphas_cumprod_prev F.pad(alphas_cumprod[:-1], (1, 0), value1.0) # ᾱ_{t-1} 且 ᾱ_0 1 # 用于从 x_0 和 ε 计算 x_t 的系数 sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1. - alphas_cumprod) # 用于从 x_t 和 ε_θ 计算 x_{t-1} 的系数参见反向过程公式 # 计算 x_{t-1} 的均值系数 posterior_variance betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) sqrt_recip_alphas torch.sqrt(1.0 / alphas) coefficient_for_x0 betas / torch.sqrt(1. - alphas_cumprod) # 公式中 (1 - α_t)/sqrt(1 - ᾱ_t) 的分子部分将这些系数预先计算并存储起来在训练和采样时直接查表使用避免了大量重复计算。4.3 训练与采样流程的封装一个完整的调度器类应该封装以下功能添加噪声给定x_0和随机时间步t根据公式x_t sqrt(ᾱ_t) * x_0 sqrt(1 - ᾱ_t) * ε生成加噪样本。采样步骤给定模型预测的噪声ε_θ、当前噪声图像x_t和时间步t根据反向过程公式计算x_{t-1}。这里需要处理方差σ_t的选择原始DDPM论文和其改进版DDIM有不同的选择。DDIM采样可选但重要DDIM是一种更高效的采样算法它允许用比训练步数T少得多的步数进行采样且质量下降不多。实现DDIM采样需要调整上述采样公式。注意事项数值稳定性在计算sqrt_one_minus_alphas_cumprod这类系数时特别是当alphas_cumprod非常接近1时可能会遇到数值下溢问题。一个实用的技巧是使用torch.clamp将其限制在一个很小的正数以上例如torch.sqrt(1. - torch.clamp(alphas_cumprod, min1e-8))。5. 实战演练训练循环与图像生成有了UNet和调度器我们就可以将它们组合起来完成模型的训练和推理。5.1 训练循环的编写DDPM的训练出奇地简洁和稳定其核心步骤如下def train_one_epoch(model, scheduler, dataloader, optimizer, device): model.train() total_loss 0 for batch in dataloader: # 1. 准备数据 clean_images batch.to(device) # (B, C, H, W) batch_size clean_images.shape[0] # 2. 随机采样时间步 t torch.randint(0, scheduler.timesteps, (batch_size,), devicedevice).long() # 3. 根据前向过程公式添加噪声 noise torch.randn_like(clean_images) noisy_images scheduler.add_noise(clean_images, t, noise) # 4. 模型预测噪声 predicted_noise model(noisy_images, t) # 5. 计算损失预测噪声与真实噪声的均方误差 loss F.mse_loss(predicted_noise, noise) # 6. 反向传播与优化 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader)为什么训练如此简单关键在于损失函数。它直接比较预测噪声和真实噪声的像素级MSE。这个目标函数被证明是优化数据对数似然变分下界ELBO的一个有效简化。在实践中它非常鲁棒不容易出现GAN训练中常见的模式崩溃问题。5.2 采样图像生成过程采样是训练的反向过程我们从纯噪声开始逐步去噪。torch.no_grad() def sample(model, scheduler, image_size, batch_size4, channels3, devicecuda): # 1. 初始化为随机噪声 img torch.randn((batch_size, channels, image_size, image_size), devicedevice) # 2. 从tT逐步迭代到t1 for t in reversed(range(scheduler.timesteps)): # 当前时间步的张量形状为 (batch_size, 1) t_batch torch.full((batch_size,), t, devicedevice, dtypetorch.long) # 3. 模型预测当前步的噪声 predicted_noise model(img, t_batch) # 4. 使用调度器根据预测噪声计算更清晰的图像 (x_{t-1}) img scheduler.step(img, t_batch, predicted_noise) # 可选在中间步骤保存图像以观察生成过程 # if t % 100 0: # save_image(...) # 5. 循环结束后img即为生成的图像 x_0 # 通常需要将值从模型输出的范围如[-1,1]转换到[0,1]以便显示 img (img.clamp(-1, 1) 1) / 2.0 return img5.3 调试与可视化让训练过程“可见”在训练初期生成的结果可能完全是噪声。为了监控训练进度有几个实用的技巧固定噪声种子采样在验证集上固定一个随机噪声种子和一组时间步每隔一定训练轮数epoch就用当前的模型参数生成一次图像。通过对比这些图像你可以清晰地看到模型生成能力是如何逐步提升的——从模糊的色块到可辨认的轮廓再到清晰的细节。损失曲线监控DDPM的训练损失MSE通常会稳步下降然后在一个值附近震荡。如果损失出现NaN或剧烈上升很可能是学习率过高、数值不稳定或数据预处理有问题。反向过程可视化在采样函数中将每一步的img保存下来可以制作成一个视频。你会看到图像如何从雪花点般的噪声中“逐渐浮现”出来这是一个非常直观的理解扩散过程的方式。踩坑实录硬件限制与调优在单张消费级GPU如RTX 3080 10GB上训练一个64x64分辨率的DDPM模型是可行的。但如果将分辨率提升到128x128或更高UNet的参数量和中间激活值会急剧增加很容易导致显存溢出OOM。此时必须采用梯度检查点Gradient Checkpointing技术它用计算时间换显存空间能显著降低训练所需显存。在PyTorch中可以使用torch.utils.checkpoint.checkpoint函数包装UNet中的某些模块。另一个技巧是使用混合精度训练AMP它既能节省显存又能加速训练。 从零实现DDPM：PyTorch复现扩散模型核心原理与代码实践
1. 从噪声到图像DDPM的核心思想与复现价值如果你对生成式AI感兴趣那么“扩散模型”这个词你一定不陌生。从DALL-E 2到Stable Diffusion这些能根据文字描述生成逼真图像的“魔法”背后核心引擎就是扩散模型。而DDPM即去噪扩散概率模型可以说是现代扩散模型的奠基之作。它用一套极其优雅的数学框架教会了神经网络如何从纯粹的随机噪声中一步步“雕刻”出我们想要的图像。网上关于DDPM原理的解读很多公式推导也足够深入但当你真正想动手实现一个时往往会发现理论与代码之间存在一道鸿沟。那些论文里的“前向过程”、“反向过程”、“变分下界”如何用PyTorch的Tensor操作表达出来时间步长t该如何嵌入到网络里那个关键的UNet结构到底长什么样每一层为什么要这么设计扩散调度中的β_t序列又该如何影响训练和采样这篇文章我将带你从零开始用PyTorch完整复现一个DDPM。我们不只停留在调用高级API而是要亲手搭建时间嵌入模块、构建UNet网络、实现扩散调度器并理解其中每一个超参数和设计选择的用意。这个过程就像组装一台精密的钟表你会看清每一个齿轮模块是如何咬合最终让整个系统运转起来。无论你是想深入理解扩散模型的内在机理还是为后续研究更复杂的模型如带条件的扩散模型打下坚实基础这次“造轮子”的经历都将让你获益匪浅。2. 扩散模型的“剧本”前向与反向过程详解在搭建任何模型之前我们必须彻底理解它要解决什么问题以及解决问题的“剧本”是什么。DDPM的剧本分为截然相反的两幕破坏与重建。2.1 前向过程给图像逐步加噪第一幕是前向过程或者叫扩散过程。这个过程是确定的、无需学习的。想象我们有一张清晰的原始图片x_0。我们准备了一个包含T个时间步的“噪声添加计划”。在每一个时间步t我们向图片中添加一点点高斯噪声。具体来说我们从x_{t-1}得到x_t的公式是x_t sqrt(1 - β_t) * x_{t-1} sqrt(β_t) * ε_t其中ε_t是从标准正态分布中采样的噪声β_t是一个预先定义好的、很小的正数例如从0.0001线性增长到0.02它决定了每一步添加的噪声量。sqrt(1 - β_t)和sqrt(β_t)是为了保证每一步的方差得到合理控制。这个过程有一个非常美妙的性质由于每一步都是线性高斯变换我们可以通过重参数化技巧直接从原始图像x_0计算出任意中间时刻t的加噪图像x_t而无需一步步迭代。公式如下x_t sqrt(ᾱ_t) * x_0 sqrt(1 - ᾱ_t) * ε这里α_t 1 - β_t而ᾱ_t Π_{s1}^{t} α_s即所有α_s的连乘。ε同样是标准高斯噪声。这个公式是代码实现的关键它意味着我们可以在训练时随机采样一个时间步t然后直接用这个公式构造出对应的加噪样本x_t效率极高。前向过程的本质它是一个将数据分布清晰的图片逐渐转变为简单先验分布标准高斯噪声的过程。当T足够大时x_T就几乎完全是一团纯粹的随机噪声了。2.2 反向过程让神经网络学会去噪第二幕是反向过程这才是模型需要学习的部分。我们的目标是逆转前向过程给定一个噪声图像x_t和时间步t我们想要预测出那一步所添加的噪声ε或者等价地预测出更清晰的上一时刻图像x_{t-1}。DDPM选择让神经网络ε_θ去预测噪声。为什么是预测噪声而不是直接预测x_{t-1}呢这背后有数学上的考量与变分下界的推导有关但从直觉上理解也很直接预测一个均值为0、方差为1的标准高斯噪声分量对于神经网络来说是一个更稳定、更容易学习的任务。因此反向过程的核心步骤是向神经网络输入加噪图像x_t和时间步t。神经网络输出对噪声ε的预测ε_θ(x_t, t)。利用预测的噪声我们可以根据以下公式估算出x_{t-1}x_{t-1} 1 / sqrt(α_t) * (x_t - (1 - α_t)/sqrt(1 - ᾱ_t) * ε_θ(x_t, t)) σ_t * z其中z是额外的随机噪声当t1时加入以增加生成样本的多样性σ_t是其方差。通过从tT到t1迭代执行上述步骤我们就能从一团噪声x_T ~ N(0, I)开始逐步“去噪”最终得到一张清晰的图像x_0。理解这个“剧本”的价值它决定了我们代码的整体框架。我们的模型UNet扮演的是“噪声预测器”的角色而训练目标就是最小化预测噪声ε_θ与真实添加噪声ε之间的差距。接下来我们就来搭建这位核心演员——UNet。3. 构建核心演员时间感知的UNet网络UNet是DDPM中用于噪声预测的骨干网络其结构像一个“U”形先下采样编码提取特征再上采样解码重建图像。但DDPM中的UNet有一个关键变体它必须能感知“时间”。因为去噪的动作强烈依赖于当前处于去噪过程的哪一步时间步t网络需要知道“我现在是在清理早期的大块噪声还是在做最后的精修”。3.1 时间步嵌入让网络拥有“时间感”时间步t是一个标量整数如250我们需要将其转化为网络能够利用的特征。直接输入全连接层是一种方法但DDPM论文采用了Transformer中类似的正弦位置编码思路并将其通过一个小的MLP多层感知机进行变换。import math import torch import torch.nn as nn class TimeEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim dim # 构建一个包含两个线性层和SiLU激活的小型MLP self.mlp nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), # SiLU (Swish)激活函数实践中比ReLU在扩散模型中效果更好 nn.Linear(dim * 4, dim) ) def forward(self, t): # t: (batch_size, 1) 整数时间步 half_dim self.dim // 2 # 生成频率因子和Transformer的位置编码类似 freqs torch.exp( -math.log(10000) * torch.arange(start0, endhalf_dim, dtypetorch.float32) / half_dim ).to(t.device) # 将t与频率因子相乘得到角度参数 args t.float() * freqs # 分别计算正弦和余弦部分然后拼接得到位置编码 embedding torch.cat([torch.sin(args), torch.cos(args)], dim-1) # 如果嵌入维度是奇数进行填充以匹配维度 if self.dim % 2: embedding torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim-1) # 通过MLP进行非线性变换得到最终的时间嵌入向量 return self.mlp(embedding)这个TimeEmbedding模块的输出是一个(batch_size, dim)的张量。接下来我们需要将这个时间信息注入到UNet的每一层中。3.2 UNet的层结构残差块与注意力机制一个标准的DDPM UNet由多个“下采样块”和“上采样块”组成。每个块的核心是一个“残差块”它负责主要的特征变换。此外在网络的深层特征图尺寸较小处会插入“注意力块”让模型能够建立图像不同区域之间的长程依赖这对于生成结构连贯的图像至关重要。残差块的设计它需要融合三种信息输入特征图x、时间嵌入t_emb以及在某些情况下对于条件生成的条件信息。典型的做法是对输入特征x进行GroupNorm和SiLU激活。将时间嵌入t_emb通过一个线性层投影到与特征图通道数匹配的维度然后加到特征上通常是在某个归一化层之后。使用一个卷积层进行特征变换。如果该块有残差连接即输入输出通道数相同则将变换后的特征与输入相加。注意力块的引入在特征图尺寸为16x16或8x8的层级我们会使用多头自注意力机制。由于此时空间维度已大大降低计算注意力权重的开销是可接受的。注意力机制能让模型在去噪时更好地协调图像不同部分的内容例如确保生成的猫有两只对称的眼睛。3.3 完整的UNet组装与下/上采样UNet的整体流程如下初始卷积将输入的(batch, 3, H, W)图像映射到初始通道数如128。编码器路径下采样由多个阶段组成。每个阶段包含若干个残差块每个块都注入当前时间步信息。可能的一个注意力块在深层。一个下采样层通常用步长为2的卷积或池化将特征图尺寸减半通道数翻倍。瓶颈层在最低分辨率下使用多个连续的残差块和注意力块进行深度处理。解码器路径上采样与编码器对称。每个阶段包含一个上采样层转置卷积或最近邻插值卷积将特征图尺寸翻倍通道数减半。若干个残差块。关键点解码器每个残差块的输入除了来自上一层的特征和当前时间嵌入还会拼接Concatenate来自编码器路径对应层级的特征。这是UNet的核心思想它让解码器在重建细节时能利用编码器提取的底层信息。最终卷积将通道数映射回3RGB并使用tanh或sigmoid激活将值约束到[-1, 1]或[0, 1]的像素范围。实操心得通道数与下采样倍数在设计UNet时一个常见的陷阱是下采样次数太多。对于64x64的图像下采样4次64-32-16-8-4是合理的。但如果下采样到1x1会丢失大量空间信息导致生成质量下降。同时初始通道数不宜过小如32通常128或256是一个好的起点以确保模型有足够的容量学习复杂的去噪映射。4. 导演的指令表扩散调度器的实现如果说UNet是演员那么扩散调度器就是导演手中的剧本和指令表。它严格定义了前向过程中每一步的噪声强度β_t并提供了反向采样时所需的全部系数。调度器的设计直接影响训练稳定性和生成质量。4.1 β_t 序列的生成线性与余弦调度最直接的策略是线性调度即让β_t从一个小值β_start线性增长到一个大值β_end。def linear_beta_schedule(timesteps, beta_start1e-4, beta_end0.02): return torch.linspace(beta_start, beta_end, timesteps)然而线性调度在过程的两端非常清晰和非常噪声变化可能过于剧烈。后续研究提出了改进的余弦调度它能让ᾱ_t即x_t中保留的原始信号系数随着t的变化遵循一个余弦函数的一半周期从而在过程开始和结束时变化更平缓。def cosine_beta_schedule(timesteps, s0.008): steps timesteps 1 x torch.linspace(0, timesteps, steps) alphas_cumprod torch.cos(((x / timesteps) s) / (1 s) * math.pi * 0.5) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] # 归一化使ᾱ_01 betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0, 0.999)在实际操作中余弦调度通常能带来更平滑的训练过程和更好的生成效果尤其是在步数较多时。4.2 关键系数的预计算一旦确定了β_t序列我们就可以预先计算出一系列在训练和采样中反复使用的系数并将其注册为模型的缓冲张量这能极大提升效率。# 假设 betas 是形状为 (T,) 的张量 alphas 1. - betas alphas_cumprod torch.cumprod(alphas, dim0) # ᾱ_t alphas_cumprod_prev F.pad(alphas_cumprod[:-1], (1, 0), value1.0) # ᾱ_{t-1} 且 ᾱ_0 1 # 用于从 x_0 和 ε 计算 x_t 的系数 sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1. - alphas_cumprod) # 用于从 x_t 和 ε_θ 计算 x_{t-1} 的系数参见反向过程公式 # 计算 x_{t-1} 的均值系数 posterior_variance betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) sqrt_recip_alphas torch.sqrt(1.0 / alphas) coefficient_for_x0 betas / torch.sqrt(1. - alphas_cumprod) # 公式中 (1 - α_t)/sqrt(1 - ᾱ_t) 的分子部分将这些系数预先计算并存储起来在训练和采样时直接查表使用避免了大量重复计算。4.3 训练与采样流程的封装一个完整的调度器类应该封装以下功能添加噪声给定x_0和随机时间步t根据公式x_t sqrt(ᾱ_t) * x_0 sqrt(1 - ᾱ_t) * ε生成加噪样本。采样步骤给定模型预测的噪声ε_θ、当前噪声图像x_t和时间步t根据反向过程公式计算x_{t-1}。这里需要处理方差σ_t的选择原始DDPM论文和其改进版DDIM有不同的选择。DDIM采样可选但重要DDIM是一种更高效的采样算法它允许用比训练步数T少得多的步数进行采样且质量下降不多。实现DDIM采样需要调整上述采样公式。注意事项数值稳定性在计算sqrt_one_minus_alphas_cumprod这类系数时特别是当alphas_cumprod非常接近1时可能会遇到数值下溢问题。一个实用的技巧是使用torch.clamp将其限制在一个很小的正数以上例如torch.sqrt(1. - torch.clamp(alphas_cumprod, min1e-8))。5. 实战演练训练循环与图像生成有了UNet和调度器我们就可以将它们组合起来完成模型的训练和推理。5.1 训练循环的编写DDPM的训练出奇地简洁和稳定其核心步骤如下def train_one_epoch(model, scheduler, dataloader, optimizer, device): model.train() total_loss 0 for batch in dataloader: # 1. 准备数据 clean_images batch.to(device) # (B, C, H, W) batch_size clean_images.shape[0] # 2. 随机采样时间步 t torch.randint(0, scheduler.timesteps, (batch_size,), devicedevice).long() # 3. 根据前向过程公式添加噪声 noise torch.randn_like(clean_images) noisy_images scheduler.add_noise(clean_images, t, noise) # 4. 模型预测噪声 predicted_noise model(noisy_images, t) # 5. 计算损失预测噪声与真实噪声的均方误差 loss F.mse_loss(predicted_noise, noise) # 6. 反向传播与优化 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader)为什么训练如此简单关键在于损失函数。它直接比较预测噪声和真实噪声的像素级MSE。这个目标函数被证明是优化数据对数似然变分下界ELBO的一个有效简化。在实践中它非常鲁棒不容易出现GAN训练中常见的模式崩溃问题。5.2 采样图像生成过程采样是训练的反向过程我们从纯噪声开始逐步去噪。torch.no_grad() def sample(model, scheduler, image_size, batch_size4, channels3, devicecuda): # 1. 初始化为随机噪声 img torch.randn((batch_size, channels, image_size, image_size), devicedevice) # 2. 从tT逐步迭代到t1 for t in reversed(range(scheduler.timesteps)): # 当前时间步的张量形状为 (batch_size, 1) t_batch torch.full((batch_size,), t, devicedevice, dtypetorch.long) # 3. 模型预测当前步的噪声 predicted_noise model(img, t_batch) # 4. 使用调度器根据预测噪声计算更清晰的图像 (x_{t-1}) img scheduler.step(img, t_batch, predicted_noise) # 可选在中间步骤保存图像以观察生成过程 # if t % 100 0: # save_image(...) # 5. 循环结束后img即为生成的图像 x_0 # 通常需要将值从模型输出的范围如[-1,1]转换到[0,1]以便显示 img (img.clamp(-1, 1) 1) / 2.0 return img5.3 调试与可视化让训练过程“可见”在训练初期生成的结果可能完全是噪声。为了监控训练进度有几个实用的技巧固定噪声种子采样在验证集上固定一个随机噪声种子和一组时间步每隔一定训练轮数epoch就用当前的模型参数生成一次图像。通过对比这些图像你可以清晰地看到模型生成能力是如何逐步提升的——从模糊的色块到可辨认的轮廓再到清晰的细节。损失曲线监控DDPM的训练损失MSE通常会稳步下降然后在一个值附近震荡。如果损失出现NaN或剧烈上升很可能是学习率过高、数值不稳定或数据预处理有问题。反向过程可视化在采样函数中将每一步的img保存下来可以制作成一个视频。你会看到图像如何从雪花点般的噪声中“逐渐浮现”出来这是一个非常直观的理解扩散过程的方式。踩坑实录硬件限制与调优在单张消费级GPU如RTX 3080 10GB上训练一个64x64分辨率的DDPM模型是可行的。但如果将分辨率提升到128x128或更高UNet的参数量和中间激活值会急剧增加很容易导致显存溢出OOM。此时必须采用梯度检查点Gradient Checkpointing技术它用计算时间换显存空间能显著降低训练所需显存。在PyTorch中可以使用torch.utils.checkpoint.checkpoint函数包装UNet中的某些模块。另一个技巧是使用混合精度训练AMP它既能节省显存又能加速训练。