去噪扩散概率模型DDPM的PyTorch实现与源码解析

简介：这是一套基于Pytorch实现的去噪扩散概率模型（DDPM）完整项目源码，面向希望深入理解扩散模型原理、并动手实践图像去噪与增强的开发者、研究者和学生。项目代码涵盖数据加载、网络构建、损失函数与优化器配置等完整流程，同时清晰呈现了前向扩散过程、反向去噪过程、训练策略以及模型评估等关键模块，能够帮助读者将深度学习理论快速落地为可运行、可调试的实战项目。资源包共22个文件，以19个Python脚本为核心，包含多种UNet变体、引导扩散、重采样等实现，并配有2张示例效果图和1份说明文档，整体压缩包仅967KB，结构紧凑、目录清晰，便于按需查阅和二次开发。已有424人浏览学习，对于入门或进阶深度学习图像处理方向具有不错的参考价值。通过运行和研究这套源码，读者可系统掌握DDPM的完整实现细节，也可基于现有框架进行改进创新，从而提升模型开发与调优的实战能力。

1. 全是噪点时别慌：去噪扩散概率模型（DDPM）的 PyTorch 落地

如果你刚接触图像生成，第一次跑通去噪扩散概率模型（DDPM）大概会经历三个阶段：前几百步输出纯噪声，以为数据路径配错了；几千步后出现模糊轮廓，开始怀疑模型过拟合了某一张图；直到几万步后采样出清晰样张，才确认它真的在学数据分布。这套基于 PyTorch 的源码包把 DDPM 完整工程化，把经典 DDPM、Karras U-Net、无分类器引导、RePaint 图像修复、FID 评估收进同一套代码结构，适合想理解原理并亲手跑通训练的初学者，也适合需要对比不同变体效果的工程研究者。从数学骨架拆起，落点始终是能直接运行的脚本和可复现的参数。

2. 前向加噪与反向去噪：扩散过程的数学骨架与代码映射

2.1 前向过程的闭式解：加噪 1000 步不需要真的迭代 1000 次

DDPM 的前向过程定义为一条马尔可夫链：每一步对图像施加少量高斯噪声。如果真去逐步迭代，训练时每张图都要跑上千次，显存和耗时都不可接受。关键在于高斯分布的叠加性质——多个高斯噪声叠加后仍是高斯分布，所以第 t 步的加噪结果可以直接用闭式公式算出：

q(x_t | x_0) = N(x_t; sqrt(ᾱ_t)·x_0, (1 - ᾱ_t)·I)

其中 ᾱ_t 是前 t 步累积噪声调度的乘积，代码里用alphas_cumprod在初始化时一次性预计算，训练时直接查表。配合重参数化技巧，生成任意时间步的加噪样本只需要一行公式：

x_t = sqrt(ᾱ_t)·x_0 + sqrt(1 - ᾱ_t)·ε，ε ~ N(0, I)

源码denoising_diffusion_pytorch.py里的q_sample就是这句话的翻译。extract的作用是从预计算的一维张量中按时间步取值，并广播到与 x 相同的形状：

def q_sample(self, x_start, t, noise=None): # 未传入噪声时自动生成标准高斯噪声 noise = default(noise, lambda: torch.randn_like(x_start)) # 对应 x_t = sqrt(alpha_bar) * x_0 + sqrt(1 - alpha_bar) * eps return ( extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise )

sqrt_alphas_cumprod保留原图信息，sqrt_one_minus_alphas_cumprod注入噪声，两个张量决定给定时间步下的信噪比。调度方式的选择直接影响训练稳定性，默认的linear调度从 β=0.0001 线性增长到 0.02，而cosine调度因为各阶段破坏更均匀，高分辨率任务里更稳。两种调度的差异对照：

2.2 反向过程与损失函数：学噪声而不是学图像

前向过程负责把数据变成噪声，反向过程则要训练一个神经网络把噪声还原。DDPM 的训练目标不是直接预测去噪后的图像，而是预测每一步被注入的噪声 ε，损失函数是预测噪声与真实噪声的均方误差：

L = E_{t, x_0, ε} [ || ε - ε_θ(x_t, t) ||² ]

这样设计的好处是目标函数稳定，不引入对抗训练，也不需要逐层重构。模型输入是加噪图像和时间步，输出是与输入同尺寸的噪声预测。训练主循环的核心代码是 GaussianDiffusion 的forward，Trainer 每步调用它就能拿到标量 loss：

def forward(self, x, *args, **kwargs): b, c, h, w, device, img_size = *x.shape, x.device, self.image_size # 每个 batch 随机抽取时间步，保证各图噪声水平不同 t = torch.randint(0, self.num_timesteps, (b,), device=device).long() x = x * 2 - 1 # 像素归一化到 [-1, 1]，与噪声量纲对齐 noise = torch.randn_like(x) x_noisy = self.q_sample(x_start=x, t=t, noise=noise) pred_noise = self.denoise_fn(x_noisy, t) loss = F.mse_loss(pred_noise, noise) return loss

torch.randint让同一个 batch 里每张图的破坏程度不同，模型在每个训练步都能见到从轻微到重度去噪的全部难度区间。x * 2 - 1这一步很多人会漏，如果数据还停留在 [0, 1] 范围，模型输出与标准高斯噪声之间始终存在偏移，loss 怎么降都降不到合理区间。

2.3 采样循环：反向过程如何一步步还原图像

训练只产出噪声预测，真正生成图像发生在采样阶段。采样从纯噪声 x_T 出发，按 t 从大到小迭代，每步用模型预测的噪声反推 x_{t-1}。关键方法p_sample的实现如下：

@torch.no_grad() def p_sample(self, x, t, t_index): # 模型预测当前步的噪声 pred_noise = self.denoise_fn(x, t) # 利用前向过程的逆公式估算原始图像 x_start = self.predict_start_from_noise(x, t, pred_noise) if t_index == 0: return x_start # 均值项来自去噪后的估计，随机项来自后验方差 model_mean = self.model_mean(x, t, x_start, pred_noise) posterior_variance = self.posterior_variance(x, t) noise = torch.randn_like(x) return model_mean + extract(posterior_variance.sqrt(), t, x.shape) * noise

predict_start_from_noise是前向公式的代数逆运算，从当前噪声图像和预测噪声解出 x_0 的估计。t_index == 0时直接返回 x_start，最后一步不需要再加随机噪声。采样质量对posterior_variance的精度敏感，这就是为什么所有调度张量都用 float32 预计算，而不是采样时临时生成。

3. 源码结构拆解：一套可以按需拼装的扩散模型工具箱

3.1 模块清单与职责边界

这套源码不是单文件脚本，而是按功能拆成多个模块。第一次拿到压缩包先别急着跑训练，把文件归归类，后面改需求会省很多时间。下面是模块与使用场景的对照：

我的做法是先读denoising_diffusion_pytorch.py，它把模型、扩散过程和训练器三个对象分离得最清楚，其他文件基本沿用同一套设计。__init__.py只导出主模块的类，变体模块必须用from denoising_diffusion_pytorch.repaint import ...的方式单独导入，这一点容易看漏。压缩包根目录的 README.md 和 images 下的样例图给出了预期输出形态，跑训练之前先看一眼能少走弯路。另外continuous_time_gaussian_diffusion.py、v_param_continuous_time_gaussian_diffusion.py和weighted_objective_gaussian_diffusion.py属于连续时间扩散与 v 参数化变体，追求少步采样或训练加速时才需要深入研究，初学者可以先跳过。

3.2 U-Net 骨干与注意力机制：attend.py 里的内存优化

扩散模型对 U-Net 骨干的要求比图像分割更高，因为每步都要在整图上做去噪，通道数与注意力机制直接决定生成质量。Unet 的核心可控参数是dim（基础通道数）和dim_mults（各下采样阶段的通道倍数）。U-Net 的另一个关键输入是时间步编码——模型必须知道当前去噪到第几步，所以每个时间步先过正弦位置编码，再映射成与通道数一致的 embedding 逐层注入，attend.py里的注意力模块负责接收这个条件。

attend.py封装了几种注意力实现，核心开关是flash_attn：

from denoising_diffusion_pytorch.attend import Attention attn = Attention( dim=512, # 输入特征维度 dim_head=64, # 每个注意力头的维度 heads=8, # 注意力头数 flash=True # 开启 Flash Attention 降低显存 )

flash=True走 Flash Attention 路径，显存占用显著下降，但需要 GPU 支持；显存不紧张时开不开对结果影响不大。batch size 上不去的场景，优先打开它比盲目减通道更划算。

3.3 不同变体的取舍：从 simple 到 karras 到 3D

simple_diffusion.py适合先把思路验证通，训练循环是最小可读的；karras_unet.py的骨干更接近 Karras 等人论文的设计，在分辨率较高、训练步数充足时收敛更稳。karras_unet_1d.py和karras_unet_3d.py处理不同维度的数据空间：一维适合音频或传感器信号，三维适合视频或医学体数据。维度不同，卷积核、下采样方式、注意力作用范围都要跟着改，骨干参数表混用会导致维度不匹配的报错。

选变体时先回答三个问题：数据是几维的、是否需要条件控制、训练资源允许多大的 batch。答案基本就定位到具体文件了。想用 Karras 骨干但沿用默认 DDPM 训练循环也完全可行，类接口是兼容的，只要 image_size 与通道数对得上。

4. 跑通训练与采样：环境配置、训练脚本与参数边界

4.1 环境准备与安装

这套代码依赖 PyTorch 和 torchvision，PyTorch 2.x 是当前常用版本。建议用 conda 建独立环境，避免和现有项目互相污染：

conda create -n ddpm python=3.10 -y conda activate ddpm pip install torch torchvision pip install -e .

pip install -e .会读取项目根目录的 setup.py 完成安装，依赖声明在 install_requires 里会自动拉取；以可编辑模式安装后，改动源码导出的永远是当前最新版本，不用重复装。安装完成后在 Python 里执行import denoising_diffusion_pytorch不报错，环境即就绪。

4.2 最小可运行训练脚本

经典主模块的完整训练脚本只需要三个对象：Unet 骨干、GaussianDiffusion 扩散过程、Trainer 训练器。下面的脚本可以直接照抄，把数据路径换成自己的图片目录：

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer model = Unet( dim=64, dim_mults=(1, 2, 4, 8), flash_attn=True ) diffusion = GaussianDiffusion( model, image_size=128, timesteps=1000, sampling_timesteps=250 ) trainer = Trainer( diffusion, './data/custom_images', train_batch_size=16, train_lr=1e-4, train_num_steps=700000, gradient_accumulate_every=2, ema_decay=0.995, amp=True ) trainer.train()

dim=64是 U-Net 最底层通道数，dim_mults决定每个下采样阶段的通道翻倍方式。image_size必须能被2 ** len(dim_mults)整除，上面配置下是 128 能被 16 整除，改成 100 会直接断言失败。timesteps=1000是论文标准配置，sampling_timesteps=250表示采样阶段从 1000 步中均匀取 250 步执行，实际走的是 DDIM 式跳步路径，生成速度快约 4 倍，质量损失很小。

4.3 训练参数调整表与判断依据

训练参数不是越大越好，关键看显存和 loss 曲线。一组常用参数和调整方向整理成表，遇到问题先查表再动手：

loss 曲线应该平滑下降，不用追求降到 0，因为目标是最小化噪声预测误差，过小的 loss 往往意味着模型退化成了预测均值。EMA 权重在采样时通常比原始权重效果好，Trainer 内部已经处理，不需要自己维护第二份权重。

4.4 采样与断点续训

Trainer 在训练过程中会自动保存最近的权重和优化器状态。采样的入口是 diffusion 对象的sample方法：

from denoising_diffusion_pytorch import Unet, GaussianDiffusion model = Unet(dim=64, dim_mults=(1, 2, 4, 8)) diffusion = GaussianDiffusion(model, image_size=128, timesteps=1000) diffusion.load_state_dict(torch.load('./results/model-100.pt', map_location='cpu')) samples = diffusion.sample(batch_size=4)

注意：sample返回 [-1, 1] 范围的张量，显示或保存前先执行(samples + 1) / 2再乘 255 转图像，直接保存会得到全黑或全灰的图，这是最常被误判成训练失败的一步。

5. 进阶玩法：引导采样、RePaint 修复与 FID 评估的注意事项

5.1 无分类器引导：一个参数把无条件模型变成条件模型

无分类器引导的思路，是在同一个模型里同时学无条件分布和条件分布。训练时给 U-Net 传入类别标签，并按cond_drop_prob概率随机把类别置空，采样时再用引导权重把噪声预测从无条件方向推向条件方向。classifier_free_guidance.py的用法如下：

from denoising_diffusion_pytorch.classifier_free_guidance import Unet, GaussianDiffusion, Trainer unet = Unet( dim=64, dim_mults=(1, 2, 4, 8), num_classes=10, cond_drop_prob=0.5 # 一半概率丢弃条件，学习无条件分支 ) diffusion = GaussianDiffusion(unet, image_size=64, timesteps=1000) trainer = Trainer(diffusion, './data/cifar10', train_batch_size=16) trainer.train() sampled = diffusion.sample(batch_size=8, classes=torch.tensor([3] * 8))

cond_drop_prob越大，无条件分支拟合越充分，但条件可控性下降，0.5 是稳妥起点。采样时classes必须与batch_size对齐，传从 0 开始的类别索引，长度不匹配会直接越界。引导强度一般从 2.0 起步，数值越大越贴近条件类别，多样性随之下降。

5.2 RePaint 修复：复用训练权重做局部重绘

repaint.py解决局部修复：给定原图和 mask，模型对未知区域重新加噪后再逐轮去噪，每轮把已知区域强制替换回原图，让修复区域与周围像素在统计上保持一致。它复用常规 DDPM 训练好的权重，不需要单独训练，成本最低的顺序仍是先跑通主模块训练。使用前确认 mask 布局为待修复区域 1、已知区域 0，尺寸与训练时的 image_size 一致，这两点不满足会直接得到空洞或错位的结果。

5.3 FID 评估的样本量与两个容易忽略的坑

fid_evaluation.py依赖torchmetrics[image]，缺少 image 组件时导入即报错。FID 对样本量非常敏感，少于 1 万张时统计方差很大，评估尽量用完整验证集或多次采样合并。Inception 网络要求 299x299 输入，评估代码内部会做缩放，不需要手动干预。对比两组模型时固定随机种子再采样，避免样本随机波动干扰判断。如果 FID 与肉眼观察矛盾，先查采样是否走了 EMA 权重，再查图像范围是否归一化到 [-1, 1]，这两处被改动的频率最高，排查顺序放在最前面。

本文还有配套的精品资源，点击获取