ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

DDPM扩散模型实战:从UNet到潜在扩散的完整指南

DDPM扩散模型实战:从UNet到潜在扩散的完整指南 1. 从一张模糊噪点到清晰图像DDPM到底在做什么第一次接触DDPMDenoising Diffusion Probabilistic Models的人大概率会被那一堆公式劝退。但如果你把它的核心思想翻译成人话其实就一句话先学会怎么把一张清晰图片一步步加噪变成纯噪声再反过来学会怎么从纯噪声一步步还原出清晰图片。听起来有点像“先把衣服揉成一团再学着把它熨平”但正是这种“先破坏再重建”的思路让扩散模型在图像生成质量上直接超越了此前的生成对抗网络GAN。我最初跑DDPM的时候用的是最经典的CIFAR-10数据集32×32的小图训练了大概两天。当时最大的感受是这东西训练极其稳定不像GAN那样动不动就模式崩溃或者判别器损失爆炸。你只需要定义好前向加噪过程然后让一个UNet去预测每一步加入的噪声损失函数就是简单的均方误差MSE。没有对抗训练没有复杂的博弈就是一个纯粹的回归问题。那为什么这么简单的思路能work关键在于前向过程的马尔可夫链设计。假设我们有原始图像 ( x_0 )前向过程会按照一个预设的噪声调度表 ( \beta_1, \beta_2, ..., \beta_T )逐步往图像里加高斯噪声。每一步都只依赖上一步的结果最终当 ( T ) 足够大时通常取1000步( x_T ) 就变成了一个标准正态分布。这个过程是固定的、无需学习的相当于给模型提供了一个“标准答案”——每一步加了多少噪声我们心里有数。反向过程才是真正要学的部分。我们希望训练一个神经网络 ( \epsilon_\theta )让它能够从 ( x_t ) 和当前时间步 ( t ) 出发预测出这一步被加入的噪声 ( \epsilon )。一旦能准确预测噪声就可以通过公式反推出 ( x_{t-1} )一步步往回走最终得到 ( x_0 )。这里有个非常巧妙的数学推导预测噪声和预测均值其实是等价的但预测噪声在实验中发现效果更好因为噪声的分布更接近标准正态网络更容易学习。提示很多人一开始会困惑为什么是预测噪声而不是直接预测图像。我个人的理解是噪声在整个过程中是独立同分布的而图像的分布随着时间步变化剧烈。让网络去拟合一个相对稳定的目标训练难度会低很多。从应用场景来看DDPM及其后续变体如DDIM、Latent Diffusion已经渗透到图像生成、图像修复、超分辨率、风格迁移等多个领域。尤其是Stable Diffusion这类潜在扩散模型把扩散过程放到潜空间里做大幅降低了计算成本让消费级显卡也能跑出高质量结果。如果你是一个刚入门的深度学习爱好者或者想从GAN转向扩散模型的从业者DDPM绝对是最值得啃下来的第一块硬骨头。2. 核心组件拆解UNet、噪声调度与时间步嵌入2.1 UNet为什么成了扩散模型的标配骨干打开任何一份DDPM的开源代码你都会看到一个结构高度相似的UNet。它最早是为医学图像分割设计的但后来大家发现它在扩散模型里出奇地好用。原因在于扩散模型需要输入和输出尺寸一致而且要在多个尺度上捕捉图像特征——UNet的编码器-解码器结构加上跳跃连接恰好完美匹配这个需求。具体来说DDPM里的UNet通常包含以下几个关键模块下采样路径Encoder由多个残差块Residual Block和注意力块组成每经过一个阶段特征图的空间尺寸减半通道数翻倍。这样做的目的是逐步提取从局部纹理到全局语义的多层次特征。中间层Bottleneck位于UNet最底部特征图尺寸最小但通道数最多。这里通常会插入自注意力机制Self-Attention让模型能够捕捉长距离依赖关系。对于32×32的CIFAR-10中间层可能只有4×4对于256×256的ImageNet中间层可能是8×8或16×16。上采样路径Decoder与编码器对称逐步恢复空间分辨率。每一层都会通过跳跃连接把编码器对应层的特征拼接过来帮助恢复细节信息。时间步嵌入Time Embedding这是扩散模型区别于普通UNet的关键。因为同一个UNet要在不同的时间步 ( t ) 下工作必须告诉它当前处于哪个噪声水平。通常采用正弦位置编码把标量 ( t ) 映射成一个高维向量然后通过MLP注入到每个残差块中。我实测下来UNet的深度和宽度对生成质量影响很大。在CIFAR-10上通道数设置为128、256、256、256基本够用但如果要生成256×256的人脸通道数至少翻倍而且注意力层的数量也要增加。显存不够的时候优先砍通道数而不是砍深度因为深度对感受野的影响更关键。2.2 噪声调度表的设计哲学噪声调度表决定了每一步加多少噪声直接影响到模型的学习难度和最终生成质量。最常用的两种调度方式是线性调度和余弦调度。线性调度最早在DDPM原论文中使用( \beta_t ) 从 ( 10^{-4} ) 线性增加到 ( 0.02 )。这个设置简单直接但后来大家发现它在低噪声阶段加噪太快导致图像信息过早丢失。余弦调度则让 ( \beta_t ) 按照余弦曲线变化在开始和结束阶段变化平缓中间阶段变化较快。实际对比下来余弦调度在相同训练步数下能获得更低的FID分数。下面是一个简单的线性调度实现你可以直接拿去用import torch def linear_beta_schedule(timesteps, beta_start1e-4, beta_end0.02): return torch.linspace(beta_start, beta_end, timesteps) 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) * torch.pi * 0.5) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0.0001, 0.9999)注意无论用哪种调度都要确保 ( \bar{\alpha}_T ) 足够小接近0这样最终时刻的图像才接近纯噪声。如果 ( \bar{\alpha}_T ) 还有0.1以上说明加噪不充分反向过程会很难学。2.3 时间步嵌入的几种实现方式时间步嵌入的作用是把离散的时间步 ( t ) 转换成连续的高维表示。最经典的做法是Transformer里的正弦位置编码import math import torch def sinusoidal_embedding(timesteps, dim): half_dim dim // 2 emb math.log(10000) / (half_dim - 1) emb torch.exp(torch.arange(half_dim) * -emb) emb timesteps[:, None] * emb[None, :] emb torch.cat([torch.sin(emb), torch.cos(emb)], dim-1) return emb后来也有工作尝试用可学习的嵌入或者FiLMFeature-wise Linear Modulation条件注入但正弦编码因为无需训练参数且泛化性好仍然是主流选择。我在实际项目里试过把正弦编码换成MLP直接学习结果在训练集上表现差不多但验证集上明显更差说明可学习嵌入容易过拟合。3. 动手实现从零搭建一个DDPM训练流程3.1 前向加噪过程的代码实现前向过程的核心公式是[ x_t \sqrt{\bar{\alpha}_t} x_0 \sqrt{1 - \bar{\alpha}_t} \epsilon ]其中 ( \bar{\alpha}t \prod{i1}^{t} \alpha_i )( \alpha_i 1 - \beta_i )( \epsilon \sim \mathcal{N}(0, I) )。这个公式允许我们直接从 ( x_0 ) 采样任意时刻的 ( x_t )无需一步步迭代极大提升了训练效率。class Diffusion: def __init__(self, timesteps1000, schedulecosine): self.timesteps timesteps if schedule linear: self.betas linear_beta_schedule(timesteps) else: self.betas cosine_beta_schedule(timesteps) self.alphas 1.0 - self.betas self.alphas_cumprod torch.cumprod(self.alphas, dim0) self.sqrt_alphas_cumprod torch.sqrt(self.alphas_cumprod) self.sqrt_one_minus_alphas_cumprod torch.sqrt(1.0 - self.alphas_cumprod) def q_sample(self, x0, t, noiseNone): if noise is None: noise torch.randn_like(x0) sqrt_alpha self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1) sqrt_one_minus self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1) return sqrt_alpha * x0 sqrt_one_minus * noise, noise这段代码里有个细节值得注意t的维度要和x0对齐。我见过不少新手在这里踩坑忘记做view操作导致广播错误。另外alphas_cumprod在数值上会越来越小到 ( t1000 ) 时可能只有 ( 10^{-5} ) 量级所以计算 ( \sqrt{1 - \bar{\alpha}_t} ) 时要注意浮点精度。3.2 训练循环与损失函数训练DDPM的伪代码简单到令人发指从数据集中采样一批 ( x_0 )为每个样本随机采样一个时间步 ( t )采样噪声 ( \epsilon )计算 ( x_t )让UNet预测噪声 ( \epsilon_\theta(x_t, t) )计算MSE损失并反向传播def train_step(model, diffusion, x0, optimizer): batch_size x0.shape[0] t torch.randint(0, diffusion.timesteps, (batch_size,), devicex0.device) xt, noise diffusion.q_sample(x0, t) predicted_noise model(xt, t) loss F.mse_loss(predicted_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item()这里有个经验之谈时间步采样策略会影响收敛速度。均匀采样是最常见的做法但有些工作提出按照重要性采样让模型更多关注那些损失较大的时间步。我在自己的实验里试过对中间时间步300-700加权FID确实有轻微提升但训练稳定性下降需要更仔细地调学习率。提示如果你发现损失下降很慢先检查UNet的输出是否真的在预测噪声而不是直接预测图像。另外初始化时把UNet最后一层卷积的权重初始化为0可以让模型一开始输出接近0的噪声预测训练更稳定。3.3 反向采样从纯噪声生成图像训练完成后生成过程就是从 ( x_T \sim \mathcal{N}(0, I) ) 出发逐步去噪[ x_{t-1} \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 z ]其中 ( z \sim \mathcal{N}(0, I) ) 当 ( t 1 )否则 ( z 0 )。( \sigma_t ) 可以取 ( \sqrt{\beta_t} ) 或 ( \sqrt{\tilde{\beta}_t} )。torch.no_grad() def p_sample(model, diffusion, xt, t): betas_t diffusion.betas[t] sqrt_one_minus diffusion.sqrt_one_minus_alphas_cumprod[t] sqrt_recip_alpha 1.0 / torch.sqrt(diffusion.alphas[t]) predicted_noise model(xt, t) model_mean sqrt_recip_alpha * (xt - betas_t / sqrt_one_minus * predicted_noise) if t 0: return model_mean else: noise torch.randn_like(xt) variance torch.sqrt(betas_t) * noise return model_mean variance实际跑采样的时候1000步全部走完大概需要几十秒取决于UNet大小和图像分辨率。如果嫌慢可以用DDIM采样把步数压缩到50-100步质量损失很小。我通常用DDIM 50步做快速预览最终出图再用完整DDPM。4. 踩坑实录DDPM训练与推理中的常见问题4.1 损失不下降或者生成全是噪声这是最常见的问题原因通常有以下几种问题现象可能原因排查方法损失始终在1.0左右模型输出层没有正确初始化检查最后一层卷积权重是否接近0损失下降但生成模糊训练不充分或UNet容量不足增加训练步数或通道数生成结果全是灰色时间步嵌入未正确注入打印时间步嵌入的均值和方差损失震荡严重学习率过大尝试1e-4或更低我印象最深的一次是UNet的时间步嵌入维度设成了64结果模型完全学不到时间信息生成出来的东西介于噪声和图像之间。后来把维度提到256问题立刻解决。所以时间步嵌入的维度不要太小至少要和UNet最浅层的通道数匹配。4.2 显存爆炸与批次大小调整DDPM的训练显存占用主要来自三个方面模型参数、激活值和优化器状态。以256×256图像、UNet通道数128为例batch size8时大约需要12GB显存。如果显存不够可以使用梯度累积把batch size8拆成4次forward每次batch size2使用混合精度训练AMP显存占用能降低30%左右减少UNet的通道数但不要低于64把图像裁剪成更小的patch训练注意梯度累积时损失要除以累积步数否则等效学习率会变大。4.3 采样速度太慢的优化思路完整DDPM采样1000步在RTX 3090上生成一张256×256图像大约需要20秒。如果要做交互式应用这个速度显然不够。除了DDIM还有几种加速方案DPM-Solver把采样步数压缩到20步以内质量几乎无损知识蒸馏训练一个学生模型一步生成图像潜在扩散在潜空间做扩散空间尺寸缩小8倍速度提升明显我在实际项目里最常用的是DPM-Solver20步就能出不错的结果比DDIM 50步还快。代码上只需要替换采样器训练好的UNet可以直接复用。4.4 条件生成与无分类器引导如果你想让模型按照类别或者文本生成图像就需要引入条件信息。最常用的技术是无分类器引导Classifier-Free Guidance训练时随机丢弃条件信息比如10%的概率推理时同时计算有条件预测和无条件预测然后做线性组合[ \epsilon \epsilon_{uncond} s \cdot (\epsilon_{cond} - \epsilon_{uncond}) ]其中 ( s ) 是引导强度通常取3-7。( s ) 越大生成结果越符合条件但多样性会下降。我一般先用 ( s3 ) 快速预览确定方向后再用 ( s7 ) 出最终图。5. 从DDPM到潜在扩散进阶方向与实用建议5.1 潜在扩散模型为什么更实用DDPM直接在像素空间操作计算量随图像分辨率平方增长。512×512的图像UNet中间层特征图仍然很大训练和推理都很吃力。潜在扩散模型LDM的思路是先用一个VAE把图像压缩到潜空间比如512×512×3压缩成64×64×4然后在潜空间做扩散。这样UNet的输入尺寸缩小了64倍计算量大幅下降。Stable Diffusion就是LDM的典型代表。它的UNet在64×64的潜空间上工作配合交叉注意力注入文本条件能在消费级显卡上生成1024×1024的高质量图像。如果你要做实际应用我强烈建议直接从LDM入手而不是死磕像素空间的DDPM。5.2 UNet模型改进的几个方向虽然UNet是标配但研究者一直在尝试改进它加入Transformer块用注意力替代部分卷积提升全局建模能力自适应归一化根据时间步动态调整归一化参数多尺度注意力在不同分辨率上做注意力兼顾细节和全局残差连接优化使用Pre-Activation或者Scaled Residual我试过把UNet中间层的自注意力换成窗口注意力显存占用降低了不少但生成质量有轻微下降。如果显存紧张这是一个可选的折中方案。5.3 训练数据与评估指标DDPM对数据量要求比较高CIFAR-10的5万张图能训出不错的结果但要想生成特定领域的高质量图像至少需要几万到几十万张。数据增强方面随机水平翻转是安全且有效的但颜色抖动要谨慎因为扩散模型对颜色分布很敏感。评估指标主要看FIDFréchet Inception Distance和ISInception Score。FID越低越好IS越高越好。但这两个指标都有局限性最好结合人眼主观评价。我通常会在训练过程中每10个epoch采样一批图像肉眼观察生成质量的变化趋势。提示FID对样本数量很敏感计算时至少要用5万张生成图像和全部真实图像否则数值波动很大。5.4 一个容易被忽略的细节EMA模型训练DDPM时使用指数移动平均EMA的模型参数进行采样效果通常比原始参数好很多。EMA相当于对训练过程中的参数做平滑能显著提升生成稳定性。实现上很简单class EMA: def __init__(self, model, decay0.9999): self.model copy.deepcopy(model) self.decay decay def update(self, model): for ema_param, param in zip(self.model.parameters(), model.parameters()): ema_param.data.mul_(self.decay).add_(param.data, alpha1 - self.decay)我对比过用EMA模型采样FID能降低10%-20%。这个技巧几乎零成本强烈建议加上。5.5 实际项目中的部署考量如果你要把DDPM部署到生产环境有几个点需要提前考虑模型量化把FP32转成FP16或INT8推理速度提升2-4倍批处理一次生成多张图像充分利用GPU并行能力缓存机制对于相同的随机种子和条件缓存生成结果降级方案当显存不足时自动切换到DDIM或更少步数我在一个图像生成API项目里用TensorRT对UNet做了量化单张512×512图像的生成时间从3秒降到了0.8秒。当然量化会带来轻微的质量损失需要根据业务需求权衡。最后分享一个我踩过的坑不要在生产环境用训练时的随机噪声种子。训练时为了可复现性会固定种子但推理时如果固定同一个种子每次生成的图像都一样用户会以为程序卡住了。正确做法是每次推理都从系统熵池取随机种子或者让用户自己指定。这个方向后续还可以往视频生成、3D生成扩展核心思路都是把扩散过程搬到对应的数据表示上。如果你已经把DDPM跑通了下一步可以试试在潜空间做扩散或者加入文本条件这些都是当前最活跃的研究方向。
返回列表