十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

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

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

去噪扩散概率模型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, noiseNone): # 未传入噪声时自动生成标准高斯噪声 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调度因为各阶段破坏更均匀高分辨率任务里更稳。两种调度的差异对照调度方式默认参数区间前向破坏特点适用场景linearβ 从 0.0001 到 0.02前中期信息丢失快通用基线、低分辨率cosineᾱ_t 余弦衰减各阶段信息保留均匀高分辨率、长训练步数2.2 反向过程与损失函数学噪声而不是学图像前向过程负责把数据变成噪声反向过程则要训练一个神经网络把噪声还原。DDPM 的训练目标不是直接预测去噪后的图像而是预测每一步被注入的噪声 ε损失函数是预测噪声与真实噪声的均方误差L E_{t, x_0, ε} [ || ε - ε_θ(x_t, t) ||² ]这样设计的好处是目标函数稳定不引入对抗训练也不需要逐层重构。模型输入是加噪图像和时间步输出是与输入同尺寸的噪声预测。训练主循环的核心代码是 GaussianDiffusion 的forwardTrainer 每步调用它就能拿到标量 lossdef 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,), devicedevice).long() x x * 2 - 1 # 像素归一化到 [-1, 1]与噪声量纲对齐 noise torch.randn_like(x) x_noisy self.q_sample(x_startx, tt, noisenoise) pred_noise self.denoise_fn(x_noisy, t) loss F.mse_loss(pred_noise, noise) return losstorch.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) * noisepredict_start_from_noise是前向公式的代数逆运算从当前噪声图像和预测噪声解出 x_0 的估计。t_index 0时直接返回 x_start最后一步不需要再加随机噪声。采样质量对posterior_variance的精度敏感这就是为什么所有调度张量都用 float32 预计算而不是采样时临时生成。3. 源码结构拆解一套可以按需拼装的扩散模型工具箱3.1 模块清单与职责边界这套源码不是单文件脚本而是按功能拆成多个模块。第一次拿到压缩包先别急着跑训练把文件归归类后面改需求会省很多时间。下面是模块与使用场景的对照文件职责适用场景denoising_diffusion_pytorch.py经典 DDPM 全流程Unet 加噪 采样 Trainer默认首选、通用图像生成simple_diffusion.py精简版实现教学调试、代码量最小karras_unet.py / karras_unet_1d.py / karras_unet_3d.py不同维度数据空间的 U-Net 骨干图像 / 音频信号 / 视频体数据guided_diffusion.py带引导的扩散实现条件生成任务classifier_free_guidance.py无分类器引导生成质量与多样性平衡repaint.py图像修复局部区域重绘elucidated_diffusion.pyEDM 连续时间框架少步数采样fid_evaluation.pyFID 指标评估训练后量化对比我的做法是先读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_attnfrom denoising_diffusion_pytorch.attend import Attention attn Attention( dim512, # 输入特征维度 dim_head64, # 每个注意力头的维度 heads8, # 注意力头数 flashTrue # 开启 Flash Attention 降低显存 )flashTrue走 Flash Attention 路径显存占用显著下降但需要 GPU 支持显存不紧张时开不开对结果影响不大。batch size 上不去的场景优先打开它比盲目减通道更划算。3.3 不同变体的取舍从 simple 到 karras 到 3Dsimple_diffusion.py适合先把思路验证通训练循环是最小可读的karras_unet.py的骨干更接近 Karras 等人论文的设计在分辨率较高、训练步数充足时收敛更稳。karras_unet_1d.py和karras_unet_3d.py处理不同维度的数据空间一维适合音频或传感器信号三维适合视频或医学体数据。维度不同卷积核、下采样方式、注意力作用范围都要跟着改骨干参数表混用会导致维度不匹配的报错。选变体时先回答三个问题数据是几维的、是否需要条件控制、训练资源允许多大的 batch。答案基本就定位到具体文件了。想用 Karras 骨干但沿用默认 DDPM 训练循环也完全可行类接口是兼容的只要 image_size 与通道数对得上。4. 跑通训练与采样环境配置、训练脚本与参数边界4.1 环境准备与安装这套代码依赖 PyTorch 和 torchvisionPyTorch 2.x 是当前常用版本。建议用 conda 建独立环境避免和现有项目互相污染conda create -n ddpm python3.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( dim64, dim_mults(1, 2, 4, 8), flash_attnTrue ) diffusion GaussianDiffusion( model, image_size128, timesteps1000, sampling_timesteps250 ) trainer Trainer( diffusion, ./data/custom_images, train_batch_size16, train_lr1e-4, train_num_steps700000, gradient_accumulate_every2, ema_decay0.995, ampTrue ) trainer.train()dim64是 U-Net 最底层通道数dim_mults决定每个下采样阶段的通道翻倍方式。image_size必须能被2 ** len(dim_mults)整除上面配置下是 128 能被 16 整除改成 100 会直接断言失败。timesteps1000是论文标准配置sampling_timesteps250表示采样阶段从 1000 步中均匀取 250 步执行实际走的是 DDIM 式跳步路径生成速度快约 4 倍质量损失很小。4.3 训练参数调整表与判断依据训练参数不是越大越好关键看显存和 loss 曲线。一组常用参数和调整方向整理成表遇到问题先查表再动手参数常用值作用调整方向train_batch_size16单步样本数显存溢出先减半gradient_accumulate_every2等效扩大 batch小显存替代大 batchtrain_lr1e-4优化器学习率loss 震荡时降到 5e-5ema_decay0.995权重滑动平均采样模糊时提高到 0.999ampTrue混合精度训练数值异常先关闭loss 曲线应该平滑下降不用追求降到 0因为目标是最小化噪声预测误差过小的 loss 往往意味着模型退化成了预测均值。EMA 权重在采样时通常比原始权重效果好Trainer 内部已经处理不需要自己维护第二份权重。4.4 采样与断点续训Trainer 在训练过程中会自动保存最近的权重和优化器状态。采样的入口是 diffusion 对象的sample方法from denoising_diffusion_pytorch import Unet, GaussianDiffusion model Unet(dim64, dim_mults(1, 2, 4, 8)) diffusion GaussianDiffusion(model, image_size128, timesteps1000) diffusion.load_state_dict(torch.load(./results/model-100.pt, map_locationcpu)) samples diffusion.sample(batch_size4)注意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( dim64, dim_mults(1, 2, 4, 8), num_classes10, cond_drop_prob0.5 # 一半概率丢弃条件学习无条件分支 ) diffusion GaussianDiffusion(unet, image_size64, timesteps1000) trainer Trainer(diffusion, ./data/cifar10, train_batch_size16) trainer.train() sampled diffusion.sample(batch_size8, classestorch.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]这两处被改动的频率最高排查顺序放在最前面。本文还有配套的精品资源点击获取
返回列表