简介面向深度学习与生成式模型研究者的DDPM可运行源码包聚焦扩散模型基础理论在PyTorch中的完整落地。压缩包共27个文件大小约22.21MB包含8个Python脚本、1个pth模型权重、MNIST的idx训练测试数据及相关gz压缩文件以及requirements依赖说明、README文档和生成结果示例图片等源码、模型与数据一次配齐可直接运行验证。目前已有225人学习下载。代码基于U-Net去噪网络完整覆盖DDPM类设计、前向加噪、训练与采样流程并在MNIST上完成实验生成结果与真实手写数字在视觉上较为接近。此外还探讨了不同网络架构对生成质量的影响给出关键超参数和复现注意事项便于理解扩散模型的实际调优过程。整个包目录结构清晰训练脚本、demo示例和说明文档分层放置初学者可按README从零搭建进阶者也能借源码快速开展消融实验或网络结构改进。1. DDPM 是什么一个 PyTorch 实现能给你带来什么扩散模型在图像生成领域的地位已经不用多说了而 DDPMDenoising Diffusion Probabilistic Models正是理解这一整套范式的起点。很多开发者想从零手写一个 DDPM但卡在数学公式和实现细节之间要么是加噪过程搞错要么是训练时损失不下降最后放弃了。这篇文章用一个可以直接运行的 PyTorch 实现把 DDPM 的完整链路拆开讲清楚前向加噪、反向去噪、训练循环、采样生成每个环节都有可复现代码和参数说明。无论是想入门生成模型还是想在已有项目里嵌入扩散模型做实验这套实现都能让你在本地 GPU 上跑出真实的生成效果。我会从最核心的设计直觉讲起把每一步的坑和调参经验一并交底保证你读完能独立训练并采样出自己的生成图像。2. DDPM 前向与反向过程从加噪公式到最小可运行框架DDPM 的核心思想不复杂但公式符号容易劝退。简单说前向过程是不断往图像上加高斯噪声直到原始信息完全被噪声淹没反向过程则训练一个神经网络学会一步步把噪声去掉恢复出原始图像。这个思路的关键在于它不是一步到位预测最终图像而是每一步只去除一小部分噪声让模型任务变得足够简单。2.1 前向加噪过程是如何设计的前向过程本质上是一个固定的马尔可夫链给定原始图像 x0在时间步 t 时得到 xt这个 xt 由 x0 加上一定强度的噪声得到。数学上写成xt sqrt(alpha_bar_t) * x0 sqrt(1 - alpha_bar_t) * epsilon其中 epsilon 是标准高斯噪声alpha_bar_t 是噪声调度参数的累积乘积。这个公式最大的好处是任何时间步 t 的加噪结果都可以直接从 x0 计算不需要逐步迭代。我一般把噪声调度定义为线性地从 beta_1 到 beta_T 递增这种做法简单且在 CIFAR-10 这类数据集上效果稳定。def linear_beta_schedule(num_timesteps, beta_start0.0001, beta_end0.02): return torch.linspace(beta_start, beta_end, num_timesteps) betas linear_beta_schedule(num_timesteps1000) alphas 1.0 - betas alpha_bar torch.cumprod(alphas, dim0)这段代码定义了 DDPM 中最核心的调度函数。beta_start 和 beta_end 分别控制初始和最终的噪声强度常见取值是 0.0001 到 0.02这是 DDPM 原论文的默认配置。num_timesteps 取 1000 也是原论文的设定太大训练变慢太小生成质量不稳定。alpha_bar 是累积乘积用于前向采样时直接跳转到任意时间步的加噪结果。实现里这两个参数最好用 float32 存避免累积误差。前向过程的另一个重要细节是我们只需要知道从 x0 到任意 xt 的分布这个分布由 alpha_bar 决定。所以在实现中可以直接使用上述公式批量生成噪声样本这也是训练时数据加载的主要方式。2.2 反向去噪过程为什么需要一个神经网络反向过程解决的是这样一个问题给定一个带噪图像 xt 和时间步 t如何估计出当时加入的噪声 epsilonDDPM 的结论是用一个 U-Net 结构的神经网络去预测噪声。训练目标很简单就是让预测的噪声和真实噪声之间的均方误差最小。这个设定非常优雅网络不需要直接输出图像只需要输出噪声图维度不变学习难度大幅降低。采样时我们从纯噪声 xT 出发每次用网络预测的噪声估计出 xt-1逐步迭代回到 x0这个多次迭代的过程保证了每个步骤都是一个小的去噪操作。class SimpleUNet(nn.Module): def __init__(self): super().__init__() self.down1 nn.Conv2d(3, 64, 3, padding1) self.down2 nn.Conv2d(64, 128, 3, padding1, stride2) self.down3 nn.Conv2d(128, 256, 3, padding1, stride2) self.up1 nn.ConvTranspose2d(256, 128, 4, stride2, padding1) self.up2 nn.ConvTranspose2d(128, 64, 4, stride2, padding1) self.out nn.Conv2d(64, 3, 3, padding1) def forward(self, x, t): # t 是时间步这里先用正弦编码嵌入 t_embed torch.sin(t.float() / 1000 * 2 * torch.pi) x F.relu(self.down1(x)) x F.relu(self.down2(x)) x F.relu(self.down3(x)) x F.relu(self.up1(x)) x F.relu(self.up2(x)) return self.out(x)这里给的是 U-Net 的简化演示版本实际生产用的 U-Net 会加入 ResBlock、注意力机制和 SiLU 激活但核心结构就是下采样再上采样中间层捕捉全局结构输出层恢复原分辨率。t 的时间步信息在这里只做了简单的正弦编码实际中文项目里很多人会忽略 t 的作用导致生成效果不稳定。t 编码应该作为条件信息注入到每个特征层比如通过 AdaGN 或 FiLM 机制这个在本文后面会讲到。训练时网络参数约为几百万在单张 24GB 显存 GPU 上能够顺利训练但迭代速度比现代轻量模型要慢。2.3 最小可运行的前向与采样代码把上面两个部分结合起来就可以写出一个非常精简但完整的前向加噪和采样步骤。这两个函数是整个 DDPM 实现的基础后续的训练和生成都建立在这两段代码之上也是排查 bug 的第一站。def q_sample(x0, t, noise): sqrt_alpha_bar torch.sqrt(alpha_bar[t])[:, None, None, None] sqrt_one_minus_alpha_bar torch.sqrt(1 - alpha_bar[t])[:, None, None, None] return sqrt_alpha_bar * x0 sqrt_one_minus_alpha_bar * noise def p_sample(model, x_t, t): pred_noise model(x_t, t) beta_t betas[t][:, None, None, None] alpha_t alphas[t][:, None, None, None] alpha_bar_t alpha_bar[t][:, None, None, None] mean (1 / torch.sqrt(alpha_t)) * (x_t - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * pred_noise) if t[0] 0: var beta_t * (1 - alpha_bar[t - 1]) / (1 - alpha_bar_t) noise torch.randn_like(x_t) return mean torch.sqrt(var) * noise return meanp_sample 函数实现了从 xt 到 xt-1 的一步采样。mean 的计算公式对应 DDPM 论文中的后验分布均值它由当前带噪图像、预测噪声和噪声调度共同决定。当 t 等于 0 时不需要加随机噪声直接取均值即可。这个采样循环是逐时间步进行的因此完整生成一张图像需要迭代 T 次也就是 1000 次这是 DDPM 推理速度慢的主要原因。3. 准备数据集与搭建模型让代码真正跑起来的两个关键步骤有了前向和反向的基本框架接下来要把数据加载和模型搭建落实到可以训练的状态。很多人的 DDPM 训练卡住不是公式错而是数据处理和模型结构上的细节没有处理好。3.1 选择和加载合适的训练数据集不建议在开始阶段使用高位数的真实数据集CIFAR-10 是最合适的起步选择图像尺寸是 32x32分辨率不高训练快结果可视化直观而且公开可下载。如果你已经有自己的定制数据集处理流程也是一样的无非是把图像读取和缩放这一步替换掉。训练集的准备要点是把像素归一化到 [-1, 1] 区间这个操作是必须的它和噪声调度的比例有直接关系如果像素还在 [0, 1]生成的图像会整体偏亮或者偏暗归一化到 [-1, 1] 后采样时最后一层输出加 tanh 激活即可对齐。transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) dataloader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4)这里 batch_size 取 64 是平衡显存和收敛速度的常见选择DDPM 这类扩散模型对 batch size 比较敏感太小则梯度噪声大训练不稳定。num_workers 设为 4可以让数据加载不阻塞 GPU 计算。Normalize 里的均值 0.5 和标准差 0.5 是配合 ToTensor 取值 [0,1] 的情况归一化后数据范围正好是 [-1,1]和 q_sample 公式里的噪声尺度对齐。3.2 完整模型结构U-Net 加时间嵌入前面的 SimpleUNet 只够验证流程真正要生成高质量的图片必须使用带时间嵌入和 ResBlock 的 U-Net。这里给出一个折中的 mid-level 版本它比完整版小但已经具备实际生成能力。网上流传的一些精简实现会直接把 t 拼接在通道维度上效果其实是不如条件嵌入的。class TimestepEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim dim self.fc1 nn.Linear(dim, dim * 4) self.fc2 nn.Linear(dim * 4, dim) def forward(self, t): half_dim self.dim // 2 emb math.log(10000) / (half_dim - 1) emb torch.exp(torch.arange(half_dim, devicet.device) * -emb) emb t.float().unsqueeze(1) * emb.unsqueeze(0) emb torch.cat([torch.sin(emb), torch.cos(emb)], dim-1) return F.silu(self.fc2(F.silu(self.fc1(emb))))这段时间嵌入采用的是 transformer 中广泛使用的正弦位置编码。位置编码将 t 映射到一个高维向量再经过两层全连接和 SiLU 激活生成包含时间信息的特征。时间特征随后被注入到 UNet 的每一个分辨率层级常见的方法是加到残差块的输出上。一个容易踩坑的细节是嵌入维度和特征图通道数如果不匹配会导致维度不对齐我一般把嵌入维度设计为模型最大通道数比如 256。完整的模型结构会包含多个 ResBlock 和下采样模块每层通道数通常从 64 起经过多个 stage 翻倍到 256 或 512。在下采样阶段使用 attention 模块是提高生成质量的常见做法特别是在特征图分辨率降到 16x16 和 8x8 时加入 attention 能让不同位置的特征相互关联生成图像的全局一致性会更好。过于简单的模型往往因为缺少 attention 导致生成图像出现局部错乱。3.3 优化器和损失函数的正确选择训练 DDPM 使用的是最基础的 Adam 优化器和均方误差损失。不要追求复杂的优化策略因为这个模型本身的梯度传播路径清晰使用 Adam 配合线性 warmup 就能顺利收敛。损失函数上直接在噪声预测的均方误差上优化即可。一些读者可能看到 DDPM 的另一种等价形式直接预测 x0 或者预测 v但作为实现来说预测噪声是最稳妥的。optimizer torch.optim.Adam(model.parameters(), lr2e-4) criterion nn.MSELoss() for epoch in range(num_epochs): for batch in dataloader: x0 batch[0].cuda() t torch.randint(0, num_timesteps, (x0.size(0),), devicex0.device).long() noise torch.randn_like(x0) xt q_sample(x0, t, noise) pred_noise model(xt, t) loss criterion(pred_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step()这是 DDPM 训练的完整核心循环。t 是每个样本独立随机采样的也就是说同一个 batch 里的不同图片可能处于不同的噪声阶段这一策略能让模型看到各个噪声程度的样本泛化性能更好。num_timesteps 设置越大每个样本的数据利用率越高但计算成本也会上升。训练时 loss 会从最初的 0.5 左右逐步下降到 0.1 以下如果 loss 一直在 0.5 以上不收敛说明模型结构或者学习率有问题这个可以作为排查参考。4. 训练策略与采样生成从模型收敛到出图的全流程DDPM 训练是一个相对漫长的过程但并不是完全没有技巧。这一章把训练到采样的完整流程串起来包括模型保存、采样循环和生成效果评估。4.1 训练过程中的关键监控指标只盯着 loss 曲线是不够的。在训练过程中每隔一定步数从验证集取一批图片直接做完整的前向加噪再保存中间去噪结果这是判断模型学习情况的最直接方式。如果模型学到东西你会发现从带噪图片到逐步去噪的中间结果轮廓在逐渐清晰。我一般会在每 1000 步保存一组中间可视化结果。另一个实用的监控手段是计算一个验证集上的平均 loss 作为 benchmark。训练 loss 在下降但验证 loss 不降通常意味着过拟合DDPM 中常用的缓解做法是增加训练数据增强或者降低模型容量这里随机水平翻转和随机裁剪都可以用。在扩展数据集上训练时轻度的数据增强能提升生成样本的多样性减少重复纹理。4.2 完整采样循环从纯噪声到生成图像采样是 DDPM 的推理阶段它是反向过程的实例化。前面的 p_sample 已经实现了单步去噪现在把它封装成一个完整的循环。生成一张图像需要从标准正态分布采样一个初始噪声然后迭代 T 次得到最终结果。这里的核心参数是 T 本身以及是否需要采样中间过程的可视化结果。torch.no_grad() def sample_images(model, num_images16, devicecuda): model.eval() x_t torch.randn(num_images, 3, 32, 32, devicedevice) for t in reversed(range(num_timesteps)): t_tensor torch.full((num_images,), t, devicedevice, dtypetorch.long) x_t p_sample(model, x_t, t_tensor) return x_t images sample_images(model)这段采样代码简单但完整。每一步都对整个 batch 的噪声图做一次去噪因为 p_sample 本身就是批量设计的。如果显存有限导致一次采 16 张图片放不下可以减少 num_images但要注意采样结果受初始噪声随机性的影响较大样本数太少时不容易评估生成质量。采样阶段不需要梯度使用 torch.no_grad() 可以大幅减少内存占用。模型在采样前必须切到 eval 模式否则 BatchNorm 层的行为会不同导致生成结果异常。4.3 生成质量评估视觉观察加上量化指标没有量化指标的生成实验等于白做。最简单的量化方式是用生成图像的 Fréchet Inception Distance也就是 FID 分数。这个指标越低表示生成分布和真实分布越接近。FID 计算需要 InceptionV3 提取特征常用 torchmetrics 库来做。CIFAR-10 上训练充分的 DDPM 模型 FID 可以做到十几到二十几如果训练不足会到一百以上。FID 计算在 50000 张生成样本上才是最准确的但实际调试阶段可以用 5000 张样本得到一个参考值。from torchmetrics.image.fid import FrechetInceptionDistance fid FrechetInceptionDistance(feature2048) fake_images (images 1) / 2 fid.update(real_images, realTrue) fid.update(fake_images, realFalse) print(fFID: {fid.compute().item():.2f})这一段修复了采样图像的数值范围。模型 tanh 输出的范围是 [-1, 1]而 FID 指标基于 ImageNet 预训练的分类网络输入需要归一化到 [0, 1]所以必须做一次 (x1)/2 的映射。很多实现忽略了这一步得到的 FID 指标全都异常偏高这是一个非常容易犯的隐性错误。torchmetrics 的 FID 默认使用 Inception 特征feature 参数可以设 64 或 2048经验上 2048 维特征更稳定。5. DDPM 训练常见问题与避坑指南五个最容易翻车的细节从我自己跑通 DDPM 和帮别人排查代码的经历看训练失败的现象高度集中在几个反复出现的问题上。这一部分把这些坑按现象、原因、解决方案整理出来希望能帮读者少走弯路。5.1 生成图像全黑或全白模型完全没学到东西这个问题我遇到太多次了表现是采样出来的图像是一张纯噪声图或者是固定颜色的图像。原因通常是前向加噪公式写错了。有个常见写法是把 sqrt(1 - alpha_bar) 写成 sqrt(alpha_bar)导致加噪强度计算完全错误。另一个原因是训练 loss 根本没有下降网络参数没有有效更新。解决办法是先用单张图片做前向加噪验证把 t 从 0 到 1000 的所有结果画出来肉眼确认最后一步确实变成纯噪声。然后跑一轮训练观察 loss 是否从初始值开始下降。如果 loss 完全不动检查学习率是否太低、模型是不是没有加载到 GPU、输入数据是否全为常数。我一般建议新跑一个代码库时先用 batch size 32 训练 500 步如果 loss 能明显下降再进入全量训练。5.2 训练很慢每个 batch 都要好几秒DDPM 训练速度慢是固有的因为每个 batch 都要计算从 x0 到 xt 的加噪过程而且模型本身比较大。但有些慢是代码写法问题导致的。最常见的是在 q_sample 里反复计算 sqrt_alpha_bar 这些常量由于它们不依赖数据内容完全应该在训练开始前一次性预计算并放在 GPU 上。如果在数据循环里重复计算每次会多出几十毫秒的 CPU 到 GPU 的传输开销。sqrt_alpha_bar torch.sqrt(alpha_bar).cuda() sqrt_one_minus_alpha_bar torch.sqrt(1 - alpha_bar).cuda()更进一步的优化方式是把这些常量转换成 buffer 注册到模型里也就是使用 register_buffer 让它们随模型一起移动。另外不要在前向加噪时使用 Python 循环对 t 做逐元素处理要用张量索引一次得到整个 batch 的加噪结果。这样简单的改动通常能让训练速度提升一倍以上。5.3 生成的图像虽然看起来正常但颜色偏灰或对比度低这是一个容易被忽略的问题。现象是图像的轮廓正确但颜色普遍不够鲜艳。原因大多出在最后的图像输出范围映射上。模型的输出经过 tanh 后是 [-1, 1]保存图像时需要映射到 [0, 255]。如果忘了做 (x 1) / 2 这一步直接用 x 保存就会导致图像整体偏暗且对比度低。另一个可能的原因是训练数据归一化时用的均值和标准差不对。如果你的数据集不是 CIFAR-10而是自定义数据不要套用 0.5 的均值和 0.5 的标准差。应该按数据集的实际均值和标准差来计算归一化参数否则模型学到的数据分布是有偏的生成效果自然也不会对。处理这个问题时我习惯在采样代码里写一个 save_images 函数专门做数值范围的转换和裁剪这样能避免不同环节的重复转换出错。5.4 采样时显存溢出或者采样一张图要几分钟DDPM 的采样需要 T 次前向传播1000 次推理确实是它的痛点。显存溢出的原因通常是 batch size 太大或者模型太大。如果只是做效果验证采样 batch size 设成 4 到 8 就够了。如果确实要大批量生成可以把 batch 拆成多个更小的子 batch分别采样后再合并结果。采样慢的解决方案是使用更少的采样步数。DDPM 原论文需要 1000 步但后面很多工作提出了加速方法。简单的一种是跳过中间时间步比如从 1000 步均匀采样到 100 步虽然生成质量会有所下降但在很多场景里完全可以接受。更正式的加速方法是使用 DDIM 采样器这种采样器可以在 50 步以内得到接近原版的效果。如果你想在生产环境使用 DDPM建议把 DDIM 作为标准采样器。5.5 loss 已经降下来了但采样结果还不如随机噪声这是最让人崩溃的情况。loss 明明下降得很好训练曲线也非常漂亮但采样出来的图片完全看不出任何结构。我踩过这个坑后仔细排查发现最主要的原因是时间步 t 的用法错误。如果在模型的 forward 函数里没有把 t 编码成有效特征或者说 t 被当作普通数值直接拼在特征图后面模型就无法感知当前处于哪个噪声阶段导致它在所有时间步都输出同样的预测去噪过程完全失效。另一个原因是模型结构太简单直接输出噪声图而没有任何中间层的学习容量这样的网络在 CIFAR-10 上可能能拟合训练集但生成分布和真实分布差异很大。解决方式是增加 ResBlock 数量和通道数并加入 attention 层。我用的一个经验法则是模型的参数量不能少于 2000 万否则生成质量上限会很低。6. 采样加速与进阶验证DDIM 和生成效果自查清单当 DDPM 训练正常、能生成像样的图像后接下来要解决的就是推理效率问题。DDPM 的原版采样过程太慢在真实应用中几乎不可用。DDIM 是正式解决这个问题的主流方案同时它也是理解后续各种扩散模型加速器的基础。本节的实现思路并不是什么黑匣子本质上只改动了采样公式。DDIM 的采样公式可以写成x_{t-1} sqrt(alpha_bar_{t-1}) * x0_pred sqrt(1 - alpha_bar_{t-1} - eta^2 * sigma_t^2) * epsilon_pred eta * sigma_t * noise当 eta 为 0 时这个采样过程退化为确定性采样意味着给定同一个初始噪声生成的图像是唯一确定的这在很多业务场景里非常有用。eta 大于 0 时则保留了一定的随机性用来调节生成多样性。实现 DDIM 的关键是用网络预测的噪声先反推出 x0_pred再重新组合出下一步的样本而不是直接使用原版 DDPM 的后验均值公式。torch.no_grad() def ddim_sample(model, x_t, t, t_prev): pred_noise model(x_t, t) alpha_bar_t alpha_bar[t][:, None, None, None] alpha_bar_prev alpha_bar[t_prev][:, None, None, None] x0_pred (x_t - torch.sqrt(1 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_bar_t) eps_coeff torch.sqrt(1 - alpha_bar_prev) mean torch.sqrt(alpha_bar_prev) * x0_pred eps_coeff * pred_noise return mean这段代码展示了 eta 为 0 时的确定性 DDIM 采样。整个过程只需要对预定义的时间步子集迭代不需要逐步经过每一个时间步。通过这种加速采样方式1000 步可以缩短到 50 步而生成质量不会发生明显下降。在使用时t_prev 是上一个采样时间步例如从 1000 步均匀取 50 个子步的情况下t 和 t_prev 的差值约为 20。试验中当步数低于 20 时生成质量会明显退化因此建议把采样步数控制在 20 到 100 之间。在设置子步序列时一般取 reversed 排列从噪声最大的一步开始逐步去噪这也符合逆向马尔可夫过程的自然顺序。timesteps torch.linspace(0, num_timesteps - 1, 50).long().flip(0)生成质量的自查方面我每次训练完一个新模型都会按固定顺序做三件事一是生成一张从 xT 到 x0 的逐 50 步的去噪过程图确认中间结果是一步步变清晰的这个过程可以直观看到模型在每个噪声阶段的表现是否稳定二是在固定初始噪声下采样 4 张图确认 DDIM 的确定性采样特性同一噪声输入应该得到完全一致的输出三是用 FID 在 10000 张生成样本上算一次分数确认量化指标和视觉观察一致。这三个步骤就是一套可复现的验收流程足以判断模型是否真正训练到位。实际开发中还有一个容易被忽略的小技巧把采样函数用 torch.no_grad() 包起来并且把模型切到 eval 模式这不仅能省显存还能避免 BatchNorm 对样本数量敏感造成的结果漂移。如果要在代码里批量生成样本也能用 torch.cuda.amp.autocast() 把采样过程切到半精度这在我本地的 2080Ti 上大约能带来两倍的采样加速代价是 FID 会略有波动。这个取舍取决于你的实际项目对精度的敏感度。说实话DDPM 的 PyTorch 实现并不是一个难啃的骨头但它对细节的要求比其他模型高很多。我最初跑通时也翻过车问题不是出在公式推导上而是败在了那些看起来不起眼的数值范围处理和时间步注入上。如果你正在调试自己的实现我建议先跑通本章节的完整代码把默认参数跑出结果再去动调度函数和网络结构这样能少走很多弯路。希望你也能顺利跑通自己的 DDPM有问题再回头对照这篇笔记的避坑列表排查应该能省下不少时间。本文还有配套的精品资源点击获取