简介这是一份基于扩散模型的图像恢复完整可运行代码面向深度学习与计算机视觉方向的研究者、学生及算法工程师覆盖去雨、去雾、去雪等多个典型恢复任务用户仅需调整数据集路径即可加载预训练权重直接训练或测试。资源包共三十个文件主体为十三个Python脚本负责模型定义、训练与评估、采样以及PSNR/SSIM指标计算另含YAML参数配置、XML工程配置、缓存文件和Markdown说明整体打包为29KB的RAR压缩包。目前已有14979人学习查看适合作为图像恢复任务的参考实现。代码训练和测试完整可运行附详细实验操作流程参数与数据集路径的修改方法均有说明模型、UNet骨干、扩散采样等模块分层清晰关键逻辑带注释配套博客还可进一步辅助理解便于直接迁移到自己的任务中。1. 扩散模型做图像恢复为什么大家都从 GAN 转向了 DDPM图像恢复这件事从前是 GAN 的天下。SRGAN、ESRGAN、DeblurGAN 这些名字做过超分和去噪的应该都不陌生生成器加判别器对抗训练出来的图确实锐利但代价是容易产生伪纹理——明明是一块平滑的皮肤GAN 能给你画出皮革纹路。你在测试集上 PSNR 刷得再高拿到真实 noisy 图上跑一遍翻车率不低。扩散模型Diffusion Model的出现把这个局面彻底改掉了。扩散模型做图像恢复的基本思路用一句话概括就是把还原图的过程建模成从纯噪声开始、逐步去噪的逆向过程同时把低质量图作为条件引导每一次去噪都朝着“更接近真实高清图”的方向走。它不像 GAN 那样需要判别器来“骗”生成器而是直接把每一步去噪的误差作为监督信号训练稳定得多生成结果也更少出现幻觉纹理。这个思路在去噪、去雨、超分、修复场景都能用而且同一份代码改改条件注入方式就能复现多个任务。这篇文章拆的是一份 DDPM 架构的扩散模型图像恢复完整代码附带详细实验操作流程。我按自己的复现习惯先讲清楚前向过程、逆向过程、损失函数、采样器这些关键参数背后的原理然后给出一条能直接跑通的操作链路最后把我复现时踩过的坑和排查经验列出来。不管你是刚接触扩散模型还是已经从理论视频里看得一头雾水、手里缺一份能跑的参考代码这篇笔记都能让你少走一大段弯路。2. 从 DDPM 到条件扩散模型先搞清楚去噪过程的物理意义2.1 前向过程为什么是固定添加噪声而不是学出来的DDPM 最核心的设计是前向过程完全不需要学习。给定一张高清图 $x_0$我们按照一个预设的噪声调度表noise schedule在 T 步内逐步对它添加高斯噪声。每一步的公式是# beta schedule: 线性从 1e-4 到 0.02T 步扩散 T 1000 betas torch.linspace(1e-4, 0.02, T) alphas 1.0 - betas alphas_cumprod torch.cumprod(alphas, dim0) # 累计乘积 sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1.0 - alphas_cumprod)这里alphas_cumprod是前 t 步的 α 连乘物理意义是“原始信号残留比例”。当 t 越大sqrt_alphas_cumprod越小代表原图信息被保留得越少噪声占比越多。到 t 1000 时图基本接近纯高斯噪声。这个过程是马尔可夫链每一步只依赖前一步状态而且因为高斯分布的叠加性质我们可以直接从 $x_0$ 一步采样出任意第 t 步的加噪结果不需要真的迭代 1000 次。这就是重参数化的妙处训练扩散模型不需要跑完整条马尔可夫链只需要随机抽一个 t然后让网络预测“我加的这步噪声是什么”把预测误差当作梯度。实际训练的时候t是从均匀分布里随机采样的这意味着模型在训练期间见过各个退化程度的图最终学会的是“给我一张任意噪声程度的图我能估计出从干净图到这张图所加的噪声”。这个能力就是逆向过程的基础。2.2 逆向过程和 U-Net 的条件注入低质量图从哪里进去逆向过程是扩散模型用于图像恢复的落地关键。它的目标是从随机高斯噪声 $x_T$ 出发逐步还原出 $x_0$但在每一步去噪时我们手上除了当前噪声图还有一张低质量观测图 y比如带噪图、低分辨率图或带雨纹图。这张 y 必须参与到去噪决策中否则模型只是“凭空生成一张高清图”而不是“恢复出与 y 对应的高清图”。常见做法是把 y 做 channel concat拼接到噪声图上再一起送入 U-Net。我这份代码里采用的就是这种方式# 训练时: 把低质图 y 和加噪的 x_t 在通道维拼接 x_t q_sample(x_start, t, noise) # x_start y 残差 或 x_start 高清图 x_cat torch.cat([x_t, y_cond], dim1) # 输入 U-Net 的通道数 3 3 6 pred_noise model(x_cat, t)这里有个细节你要分清。在纯 DDPM 生成任务里x_start是真实干净图条件变量不存在在图像恢复任务里通常两种建模方式一种是x_start y模型学的是从 y 的任意加噪版本还原出 y全程只优化去噪能力最后采样时以 y 为起点、只做少量去噪步另一种是x_start high_res_target即真实高清参考图低质图 y 只作为条件拼进去模型学的是“在 y 的信息约束下从噪声还原出高清版本”。我这份代码用的是第二种符合 diffusion model 图像恢复的主流范式。U-Net 的 channel 数设置我直接沿用 DDPM 默认配置base_channels 128每个 stage 翻倍attention 从 16x16 feature map 那一层开始加。训练参数量在 114M 左右单张 3090 上 batch size 8 跑 256x256 没有压力。2.3 损失函数为什么选 L1 而不是 L2扩散模型的损失函数可以选 L1 或 L2两种在数学期望上等价但 L1 的梯度在残差较大时更稳定不容易被 outlier 像素主导。我复现时先用了 L2发现训练到中后期容易出现局部震荡loss 曲线在 0.18 附近反复横跳不下降切到 L1 之后收敛明显更平滑。如果你习惯用 MSE也问题不大只是建议配合 lower learning rate1e-5 级别来压震荡。我最终使用的损失是loss F.l1_loss(model(x_cat, t), noise) # 预测噪声与真实噪声的 L1 距离这里预测的目标具体是噪声而不是直接预测图像。这是 DDPM 的原始设计——网络不学习“如何直接生成图”而是学习“如何去除已知噪声”。后文采样时根据网络预测的噪声反推出前一时刻的图像分布这个过程叫 Langevin 采样。3. 完整代码布局与训练流程一份可复现的 DDPM 图像恢复实现3.1 项目文件结构与数据准备拿到这份代码之后第一件事不是直接跑 train.py而是先把目录结构过一遍。这是我整理后的标准布局diffusion_restore/ ├── configs/ │ └── base.yaml # 训练/采样/数据全局配置 ├── data/ │ └── dataset.py # 低质图/高清图配对加载 ├── models/ │ ├── unet.py # U-Net 骨干 │ └── diffusion.py # DDPM 前向/逆向采样 ├── utils/ │ ├── metrics.py # PSNR / SSIM 计算 │ └── checkpoint.py # 模型保存与恢复 ├── train.py # 训练入口 ├── sample.py # 验证/生成入口 └── test.py # 测试集指标评估数据集部分要准备成配对格式同一张场景下一张清晰图作为 GTground truth另一张是它的退化版本——可以是自行加高斯噪声、下采样、加雨纹的合成数据。代码里定义了一个PairedDataset读取方式比较直接# data/dataset.py class PairedDataset(Dataset): def __init__(self, gt_dir, lq_dir, image_size256): self.gt_paths sorted(glob.glob(f{gt_dir}/*.png)) self.lq_paths sorted(glob.glob(f{lq_dir}/*.png)) self.image_size image_size def __getitem__(self, idx): gt Image.open(self.gt_paths[idx]).convert(RGB) lq Image.open(self.lq_paths[idx]).convert(RGB) # 统一缩放到 image_size并做随机 crop 增强 transform T.Compose([ T.RandomCrop(self.image_size), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) return transform(gt), transform(lq)归一化用 (0.5, 0.5, 0.5) 做均值和方差这会把像素值压到 [-1, 1] 区间。扩散模型里这个区间是默认操作空间因为噪声添加公式假设数据对称分布在 0 附近如果你直接用 [0, 1] 区间前向过程会在 t 较小时出现明显分布偏移恢复结果偏暗。3.2 训练入口从加载 checkpoint 到 loss 收敛判定train.py 的核心循环我做了最小化封装。先贴主训练逻辑的关键代码# train.py for epoch in range(start_epoch, epochs): for batch_idx, (gt, lq) in enumerate(train_loader): gt, lq gt.cuda(), lq.cuda() t torch.randint(0, T, (gt.size(0),), devicecuda).long() noise torch.randn_like(gt) x_t q_sample(gt, t, noise) # 对 GT 加噪 x_cat torch.cat([x_t, lq], dim1) # 和低质图拼接 pred_noise model(x_cat, t) loss F.l1_loss(pred_noise, noise) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema.update(model) # 指数滑动平均更新几个参数值得重点说明t是每个样本独立随机采的不是 batch 统一的。这样同一个 batch 内部不同样本的退化程度不同等价于隐式做数据增强避免模型只对某一段噪声水平熟练。clip_grad_norm_(1.0)很关键。扩散模型在 t 较大时噪声估计的误差梯度会异常大不加梯度裁剪我会频繁遇到 loss 突然变成 NaN。ema.update(model)使用指数滑动平均decay 我设 0.995。EMA 权重下的模型在采样阶段通常比原始权重稳定得多开头阶段可能看不出差别训练 50 epoch 后对比就很明显。训练过程中我一般每 500 步存一次 checkpoint命名带上 step 数。同时把 loss 曲线落在 tensorboard 里。我复现时用的是 AdamW初始学习率 2e-4配合 cosine schedule 衰减到 1e-6。batch size 8 对应 256x256 输入显存占用约 18GB如果你是 3090/4090 这种 24GB 卡可以正常跑显存紧的话把 batch size 降到 4学习率同步降到 1e-4不要硬撑。3.3 采样与验证DDIM 采样器如何加速推理训练完成后进入采样环节。原始 DDPM 逆向采样需要完整跑 T 步也就是 1000 次网络前向推理这在图像恢复场景里难以接受——一张图要等几十秒。代码里实现了 DDIM 采样器可以在 3050 步内完成高质量采样。核心逻辑如下# models/diffusion.py DDIM 采样 torch.no_grad() def sample_ddim(model, lq, steps50, eta0.0): model.eval() x torch.randn_like(lq) # 随机高斯噪声起点 step_seq torch.linspace(T-1, 0, steps).long() for i in range(steps-1): t step_seq[i] t_next step_seq[i1] x_cat torch.cat([x, lq], dim1) pred_noise model(x_cat, t) # DDIM 更新公式 alpha_t alphas_cumprod[t] alpha_next alphas_cumprod[t_next] x0_pred (x - torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(alpha_t) sigma eta * torch.sqrt((1 - alpha_next) / (1 - alpha_t)) * torch.sqrt(1 - alpha_t / alpha_next) c1 torch.sqrt(alpha_next) c2 torch.sqrt(1 - alpha_next - sigma**2) x c1 * x0_pred c2 * pred_noise sigma * torch.randn_like(x) return xeta0.0时 DDIM 变成确定性采样即给定同一个随机起点和同一个 lq输出严格相同这对实验对比是可贵的性质——你可以放心对比不同 checkpoint 的输出差异不会被随机性干扰。eta1.0则退化成类似 DDPM 的随机采样多样性更强但单张结果的稳定性弱一些。图像恢复任务我建议固定eta0.0既不引入额外随机性恢复质量反而更高。4. 避坑指南我复现这个项目时遇到的五个典型问题4.1 loss 曲线在 0.2 附近震荡不下去现象训练了 300 epochL1 loss 一直稳定在 0.2 附近上下振幅不超过 0.01无论调学习率还是优化器都不动。原因我排查后定位到数据归一化不一致。低质图 lq 在 dataset 里做了 Normalize(0.5)但生成 lq 图时保存的不是 [-1,1] 而是 [0,1]加载后再归一化等于做了两次偏移导致条件信息分布和网络预期不一致。解决把数据管道的归一化统一成一套GT 和 lq 都用 (0.5, 0.5, 0.5)在保存 lq 图时先把像素乘 2 减 1 再存。这之后 loss 在 30 epoch 内降到 0.15 以下。4.2 生成结果整体偏灰对比度明显不足现象采样出来的图轮廓清晰细节也能看但颜色饱和度低、暗部不够暗像蒙了一层灰纱。原因起始噪声 x_t 的采样分布和训练时不一致。我在采样代码里初始化x torch.randn_like(lq)但训练时noise torch.randn_like(gt)两者 shape 一致所以没报错可如果 lq 是经过 Resize 的它的分辨率跟 GT 不同噪声的频段分布就会有偏差。解决采样前强制把 lq resize 到和 GT 相同的尺寸再进网络。或者在数据加载阶段把 GT 和 lq 的尺寸设定为相同值不要 test 时单独用任意尺寸输入。U-Net 虽然能接受任意分辨率但扩散模型对输入分布极其敏感尺寸不匹配就是暗病。4.3 显存溢出batch size 调到 2 还是 OOM现象训练刚开始step 0 就报 CUDA out of memory直接中断。原因不是 U-Net 本身太占显存而是 attention 层在 256x256 分辨率下产生了巨大的中间特征。我的 base_channels128第四个 stage 是在 32x32 feature map 上做 attention这部分 memory 占用接近 8GB。解决把 attention 分辨率改成只在 16x16 层启用即 attention_resolution [16]会损失一点点高频细节但显存直降。另一个有效办法是开 gradient checkpointing用时间换显存代价是迭代变慢约 20%。4.4 恢复图产生重复纹理像原图被复印了好几份现象在真实低质图上测试时某些区域出现周期性重复结构尤其是纹理密集区域比如树丛、头发、草地。原因训练数据里低质图是由同一张 GT 加噪声合成的模型学会了“根据自身高频模式联想补充”也就是所谓的高频幻觉。真实退化图像的高频分布和合成噪声差异很大模型就开始了自由发挥。解决一是训练数据中混入多种退化类型不只高斯噪声加一部分下采样模糊、雨纹、JPEG 压缩噪声模型泛化性会明显提升二是推理时可以试试增加 DDIM steps 到 100更多步数意味着每一步恢复的幅度更小幻觉被逐步抑制。4.5 多卡训练 loss 降了但效果不如单卡现象用 DataParallel 跑双卡训练 loss 比单卡降得更快但最终采样效果无论是 PSNR 还是主观视觉都明显弱于单卡。原因BN 层在 DataParallel 里每个卡独立统计batch 变小导致统计量不稳定。扩散模型的 U-Net 里虽然没有显式 BN但 GroupNorm 在数据并行下有类似问题——不同卡的 normalization 统计不一致梯度更新互相干扰。解决直接用 DistributedDataParallel 替代 DataParallel每个进程独立维护完整的 batch 统计量。或者简单点放弃多卡单卡 3090 跑 300 epoch 大约 30 小时完全可接受。5. 实验流程与指标评估如何判断一组参数是否真的有效5.1 指标计算PSNR 和 SSIM 怎么算才是对的图像恢复领域最常规的客观指标是 PSNR 和 SSIM。但这里有个容易翻车的细节扩散模型输出的像素范围是 [-1,1]你和 GT 比较之前必须转回 [0,255] 的 uint8 空间如果直接在 float 张量上算 PSNR结果会整体偏高 24 dB这在论文里是绝对不允许的。这是我用的计算代码# utils/metrics.py def calculate_psnr_ssim(pred, gt, max_val255.0): pred pred.clamp(-1, 1) * 127.5 127.5 gt gt.clamp(-1, 1) * 127.5 127.5 pred pred.cpu().numpy().round().astype(np.uint8) gt gt.cpu().numpy().round().astype(np.uint8) psnr skimage.metrics.peak_signal_noise_ratio(gt, pred, data_rangemax_val) ssim skimage.metrics.structural_similarity( gt, pred, channel_axis-1, data_rangemax_val) return psnr, ssim注意structural_similarity的channel_axis参数在不同版本的 skimage 里这个位置参数有变化旧版是multichannelTrue新版是channel_axis-1。我一开始用旧参数在 scikit-image 0.22 上直接报错改成channel_axis-1就好了。PSNR 和 SSIM 不能只看绝对值要和你训练的 baseline 对比。比如你的输入 lq 本身 PSNR 是 22 dB恢复后到 28 dB这是有效提升如果只从 22 到 23.5那说明模型基本没学到东西只是在做轻度平滑。5.2 实验分组设计什么样的对比能说明模型真的有效我做扩散模型图像恢复实验时通常分成三组对比一组是输入 lq 直接双三次插值的结果作为最弱 baseline一组是用经典 U-Net 做回归预测的结果代表传统监督学习的上限一组是当前扩散模型的结果。同一份测试集上三组跑完再比较 PSNR / SSIM 和主观视觉。如果扩散模型的 PSNR 反而不如普通 U-Net 回归不要慌张。这是常见现象因为扩散模型的目标分布更接近真实高清图的流形而 PSNR 是逐像素误差GAN 和扩散模型这一类生成式方法在 PSNR 上通常拼不过直接回归方法但 SSIM 和主观效果更占优势。实际落地时我会同时报两组指标在论文或评审材料里写清楚扩散模型的优势在感知质量不在像素级误差。5.3 采样步数多少合适50 步与 250 步之间的权衡DDIM 采样步数直接决定推理耗时。我测试过同一 checkpoint 在不同 steps 下的表现10 步时图像有明显块状噪声30 步以上逐步改善50 步时和 100 步的 PSNR 差异在 0.1 dB 以内。考虑到 100 步推理耗时是 50 步的两倍实际部署我一般建议 50 步为上限对画质要求极高的场景再尝试 100 步。时间开销方面单张 3090 上50 步 DDIM 采样一张 256x256 图大约 2.5 秒50 步耗时约 1.2 秒。如果你的场景是离线批量处理这个速度可以接受如果是实时性要求高的场景得配合知识蒸馏或 latent diffusion 类方法进一步压缩步数那就超出这份代码的范围了。6. 进阶技巧用这个 DDPM 代码做超分和去雨的小改造这套代码最实用的点是它不需要伤筋动骨改动条件数据流就能迁移到超分和去雨任务。这里分享我实际改过的两条路径你拿到代码后可以照着切。6.1 从去噪切到超分只需要改数据配对和条件尺寸去噪任务的输入 lq 和输出 GT 天然同尺寸代码零改动。但超分任务里输入是下采样的小图输出是高清大图直接拼接会尺寸不匹配。我的做法是先把低分辨率图上采样到目标尺寸再拼接到噪声图上# 超分场景的输入处理 lq_resized F.interpolate(lq, size(256, 256), modebicubic, align_cornersFalse) x_cat torch.cat([x_t, lq_resized], dim1)这个做法本质上是让模型自己学会“如何补全上采样造成的模糊和缺失”。要注意modebicubic的选择不同插值方式影响模型接收到的条件信息质量。双三次是通用选择对 2x 超分效果好4x 超分时可以考虑modenearest保留边缘锐度模糊连带的信息反而更少。你可以在配置里调这个参数对比两组实验的 PSNR 差异。6.2 切到去雨任务控制合成雨纹的多样性去雨任务的退化模型是clean rain_streak雨纹的形状、方向、密度都影响模型的泛化。我合成训练数据时用了三条随机雨纹生成路径直线雨、弯曲雨、混合雨每张图随机选择 13 种叠加。重点在于不要让模型见到特定方向过硬的特征——雨纹方向最好在 [−60°, 60°] 之间随机旋转否则测试集换个方向雨纹就容易失效。训练超分和去雨时我一般把steps保持默认 1000eta0损失函数依旧 L1不需要额外改结构。唯一要调整的是数据集image_size如果你要处理 512x512 高清图显存不够就把 batch size 再降一档同时把随机 crop 改成从大图上采 256x256 的 patch 来训练推理时用滑窗拼接。这个思路在 video 领域叫 patch-based inference是验证大图效果最常用的手段。6.3 两分钟验证你的改法是否有效每次改完数据流程我都会做一个 5-epoch 的冒烟测试不跑完整训练。具体操作是在base.yaml里把epochs改成 5log_interval改成 10数据量限定在 200 张图训练完成后直接跑test.py看指标趋势python train.py --config configs/base.yaml --smoke_test python test.py --checkpoint ./checkpoints/model_step_2500.pth --test_dir ./data/test冒烟测试的关注点不是最终指标而是 loss 是否在下降、log 里的时间和显存是否符合预期以及采样输出的图像结构是否已经能看出内容轮廓。如果 5 epoch 后采样图还是一团噪声要么是学习率设置不合理要么是条件注入环节有 bug先别急着跑长训练把问题定位在源头。要说我自己的习惯从那以后我每次拿到一个新的扩散模型项目都强制自己先跑一遍冒烟测试再动完整训练。这个动作帮我挡掉了至少五六次因为数据集路径配错、归一化不一致导致的十几个小时白跑。希望这份代码和踩坑笔记也能帮你把第一个扩散模型图像恢复实验顺利跑通。本文还有配套的精品资源点击获取 扩散模型图像恢复实战：DDPM原理、代码实现与踩坑指南
简介这是一份基于扩散模型的图像恢复完整可运行代码面向深度学习与计算机视觉方向的研究者、学生及算法工程师覆盖去雨、去雾、去雪等多个典型恢复任务用户仅需调整数据集路径即可加载预训练权重直接训练或测试。资源包共三十个文件主体为十三个Python脚本负责模型定义、训练与评估、采样以及PSNR/SSIM指标计算另含YAML参数配置、XML工程配置、缓存文件和Markdown说明整体打包为29KB的RAR压缩包。目前已有14979人学习查看适合作为图像恢复任务的参考实现。代码训练和测试完整可运行附详细实验操作流程参数与数据集路径的修改方法均有说明模型、UNet骨干、扩散采样等模块分层清晰关键逻辑带注释配套博客还可进一步辅助理解便于直接迁移到自己的任务中。1. 扩散模型做图像恢复为什么大家都从 GAN 转向了 DDPM图像恢复这件事从前是 GAN 的天下。SRGAN、ESRGAN、DeblurGAN 这些名字做过超分和去噪的应该都不陌生生成器加判别器对抗训练出来的图确实锐利但代价是容易产生伪纹理——明明是一块平滑的皮肤GAN 能给你画出皮革纹路。你在测试集上 PSNR 刷得再高拿到真实 noisy 图上跑一遍翻车率不低。扩散模型Diffusion Model的出现把这个局面彻底改掉了。扩散模型做图像恢复的基本思路用一句话概括就是把还原图的过程建模成从纯噪声开始、逐步去噪的逆向过程同时把低质量图作为条件引导每一次去噪都朝着“更接近真实高清图”的方向走。它不像 GAN 那样需要判别器来“骗”生成器而是直接把每一步去噪的误差作为监督信号训练稳定得多生成结果也更少出现幻觉纹理。这个思路在去噪、去雨、超分、修复场景都能用而且同一份代码改改条件注入方式就能复现多个任务。这篇文章拆的是一份 DDPM 架构的扩散模型图像恢复完整代码附带详细实验操作流程。我按自己的复现习惯先讲清楚前向过程、逆向过程、损失函数、采样器这些关键参数背后的原理然后给出一条能直接跑通的操作链路最后把我复现时踩过的坑和排查经验列出来。不管你是刚接触扩散模型还是已经从理论视频里看得一头雾水、手里缺一份能跑的参考代码这篇笔记都能让你少走一大段弯路。2. 从 DDPM 到条件扩散模型先搞清楚去噪过程的物理意义2.1 前向过程为什么是固定添加噪声而不是学出来的DDPM 最核心的设计是前向过程完全不需要学习。给定一张高清图 $x_0$我们按照一个预设的噪声调度表noise schedule在 T 步内逐步对它添加高斯噪声。每一步的公式是# beta schedule: 线性从 1e-4 到 0.02T 步扩散 T 1000 betas torch.linspace(1e-4, 0.02, T) alphas 1.0 - betas alphas_cumprod torch.cumprod(alphas, dim0) # 累计乘积 sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1.0 - alphas_cumprod)这里alphas_cumprod是前 t 步的 α 连乘物理意义是“原始信号残留比例”。当 t 越大sqrt_alphas_cumprod越小代表原图信息被保留得越少噪声占比越多。到 t 1000 时图基本接近纯高斯噪声。这个过程是马尔可夫链每一步只依赖前一步状态而且因为高斯分布的叠加性质我们可以直接从 $x_0$ 一步采样出任意第 t 步的加噪结果不需要真的迭代 1000 次。这就是重参数化的妙处训练扩散模型不需要跑完整条马尔可夫链只需要随机抽一个 t然后让网络预测“我加的这步噪声是什么”把预测误差当作梯度。实际训练的时候t是从均匀分布里随机采样的这意味着模型在训练期间见过各个退化程度的图最终学会的是“给我一张任意噪声程度的图我能估计出从干净图到这张图所加的噪声”。这个能力就是逆向过程的基础。2.2 逆向过程和 U-Net 的条件注入低质量图从哪里进去逆向过程是扩散模型用于图像恢复的落地关键。它的目标是从随机高斯噪声 $x_T$ 出发逐步还原出 $x_0$但在每一步去噪时我们手上除了当前噪声图还有一张低质量观测图 y比如带噪图、低分辨率图或带雨纹图。这张 y 必须参与到去噪决策中否则模型只是“凭空生成一张高清图”而不是“恢复出与 y 对应的高清图”。常见做法是把 y 做 channel concat拼接到噪声图上再一起送入 U-Net。我这份代码里采用的就是这种方式# 训练时: 把低质图 y 和加噪的 x_t 在通道维拼接 x_t q_sample(x_start, t, noise) # x_start y 残差 或 x_start 高清图 x_cat torch.cat([x_t, y_cond], dim1) # 输入 U-Net 的通道数 3 3 6 pred_noise model(x_cat, t)这里有个细节你要分清。在纯 DDPM 生成任务里x_start是真实干净图条件变量不存在在图像恢复任务里通常两种建模方式一种是x_start y模型学的是从 y 的任意加噪版本还原出 y全程只优化去噪能力最后采样时以 y 为起点、只做少量去噪步另一种是x_start high_res_target即真实高清参考图低质图 y 只作为条件拼进去模型学的是“在 y 的信息约束下从噪声还原出高清版本”。我这份代码用的是第二种符合 diffusion model 图像恢复的主流范式。U-Net 的 channel 数设置我直接沿用 DDPM 默认配置base_channels 128每个 stage 翻倍attention 从 16x16 feature map 那一层开始加。训练参数量在 114M 左右单张 3090 上 batch size 8 跑 256x256 没有压力。2.3 损失函数为什么选 L1 而不是 L2扩散模型的损失函数可以选 L1 或 L2两种在数学期望上等价但 L1 的梯度在残差较大时更稳定不容易被 outlier 像素主导。我复现时先用了 L2发现训练到中后期容易出现局部震荡loss 曲线在 0.18 附近反复横跳不下降切到 L1 之后收敛明显更平滑。如果你习惯用 MSE也问题不大只是建议配合 lower learning rate1e-5 级别来压震荡。我最终使用的损失是loss F.l1_loss(model(x_cat, t), noise) # 预测噪声与真实噪声的 L1 距离这里预测的目标具体是噪声而不是直接预测图像。这是 DDPM 的原始设计——网络不学习“如何直接生成图”而是学习“如何去除已知噪声”。后文采样时根据网络预测的噪声反推出前一时刻的图像分布这个过程叫 Langevin 采样。3. 完整代码布局与训练流程一份可复现的 DDPM 图像恢复实现3.1 项目文件结构与数据准备拿到这份代码之后第一件事不是直接跑 train.py而是先把目录结构过一遍。这是我整理后的标准布局diffusion_restore/ ├── configs/ │ └── base.yaml # 训练/采样/数据全局配置 ├── data/ │ └── dataset.py # 低质图/高清图配对加载 ├── models/ │ ├── unet.py # U-Net 骨干 │ └── diffusion.py # DDPM 前向/逆向采样 ├── utils/ │ ├── metrics.py # PSNR / SSIM 计算 │ └── checkpoint.py # 模型保存与恢复 ├── train.py # 训练入口 ├── sample.py # 验证/生成入口 └── test.py # 测试集指标评估数据集部分要准备成配对格式同一张场景下一张清晰图作为 GTground truth另一张是它的退化版本——可以是自行加高斯噪声、下采样、加雨纹的合成数据。代码里定义了一个PairedDataset读取方式比较直接# data/dataset.py class PairedDataset(Dataset): def __init__(self, gt_dir, lq_dir, image_size256): self.gt_paths sorted(glob.glob(f{gt_dir}/*.png)) self.lq_paths sorted(glob.glob(f{lq_dir}/*.png)) self.image_size image_size def __getitem__(self, idx): gt Image.open(self.gt_paths[idx]).convert(RGB) lq Image.open(self.lq_paths[idx]).convert(RGB) # 统一缩放到 image_size并做随机 crop 增强 transform T.Compose([ T.RandomCrop(self.image_size), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) return transform(gt), transform(lq)归一化用 (0.5, 0.5, 0.5) 做均值和方差这会把像素值压到 [-1, 1] 区间。扩散模型里这个区间是默认操作空间因为噪声添加公式假设数据对称分布在 0 附近如果你直接用 [0, 1] 区间前向过程会在 t 较小时出现明显分布偏移恢复结果偏暗。3.2 训练入口从加载 checkpoint 到 loss 收敛判定train.py 的核心循环我做了最小化封装。先贴主训练逻辑的关键代码# train.py for epoch in range(start_epoch, epochs): for batch_idx, (gt, lq) in enumerate(train_loader): gt, lq gt.cuda(), lq.cuda() t torch.randint(0, T, (gt.size(0),), devicecuda).long() noise torch.randn_like(gt) x_t q_sample(gt, t, noise) # 对 GT 加噪 x_cat torch.cat([x_t, lq], dim1) # 和低质图拼接 pred_noise model(x_cat, t) loss F.l1_loss(pred_noise, noise) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema.update(model) # 指数滑动平均更新几个参数值得重点说明t是每个样本独立随机采的不是 batch 统一的。这样同一个 batch 内部不同样本的退化程度不同等价于隐式做数据增强避免模型只对某一段噪声水平熟练。clip_grad_norm_(1.0)很关键。扩散模型在 t 较大时噪声估计的误差梯度会异常大不加梯度裁剪我会频繁遇到 loss 突然变成 NaN。ema.update(model)使用指数滑动平均decay 我设 0.995。EMA 权重下的模型在采样阶段通常比原始权重稳定得多开头阶段可能看不出差别训练 50 epoch 后对比就很明显。训练过程中我一般每 500 步存一次 checkpoint命名带上 step 数。同时把 loss 曲线落在 tensorboard 里。我复现时用的是 AdamW初始学习率 2e-4配合 cosine schedule 衰减到 1e-6。batch size 8 对应 256x256 输入显存占用约 18GB如果你是 3090/4090 这种 24GB 卡可以正常跑显存紧的话把 batch size 降到 4学习率同步降到 1e-4不要硬撑。3.3 采样与验证DDIM 采样器如何加速推理训练完成后进入采样环节。原始 DDPM 逆向采样需要完整跑 T 步也就是 1000 次网络前向推理这在图像恢复场景里难以接受——一张图要等几十秒。代码里实现了 DDIM 采样器可以在 3050 步内完成高质量采样。核心逻辑如下# models/diffusion.py DDIM 采样 torch.no_grad() def sample_ddim(model, lq, steps50, eta0.0): model.eval() x torch.randn_like(lq) # 随机高斯噪声起点 step_seq torch.linspace(T-1, 0, steps).long() for i in range(steps-1): t step_seq[i] t_next step_seq[i1] x_cat torch.cat([x, lq], dim1) pred_noise model(x_cat, t) # DDIM 更新公式 alpha_t alphas_cumprod[t] alpha_next alphas_cumprod[t_next] x0_pred (x - torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(alpha_t) sigma eta * torch.sqrt((1 - alpha_next) / (1 - alpha_t)) * torch.sqrt(1 - alpha_t / alpha_next) c1 torch.sqrt(alpha_next) c2 torch.sqrt(1 - alpha_next - sigma**2) x c1 * x0_pred c2 * pred_noise sigma * torch.randn_like(x) return xeta0.0时 DDIM 变成确定性采样即给定同一个随机起点和同一个 lq输出严格相同这对实验对比是可贵的性质——你可以放心对比不同 checkpoint 的输出差异不会被随机性干扰。eta1.0则退化成类似 DDPM 的随机采样多样性更强但单张结果的稳定性弱一些。图像恢复任务我建议固定eta0.0既不引入额外随机性恢复质量反而更高。4. 避坑指南我复现这个项目时遇到的五个典型问题4.1 loss 曲线在 0.2 附近震荡不下去现象训练了 300 epochL1 loss 一直稳定在 0.2 附近上下振幅不超过 0.01无论调学习率还是优化器都不动。原因我排查后定位到数据归一化不一致。低质图 lq 在 dataset 里做了 Normalize(0.5)但生成 lq 图时保存的不是 [-1,1] 而是 [0,1]加载后再归一化等于做了两次偏移导致条件信息分布和网络预期不一致。解决把数据管道的归一化统一成一套GT 和 lq 都用 (0.5, 0.5, 0.5)在保存 lq 图时先把像素乘 2 减 1 再存。这之后 loss 在 30 epoch 内降到 0.15 以下。4.2 生成结果整体偏灰对比度明显不足现象采样出来的图轮廓清晰细节也能看但颜色饱和度低、暗部不够暗像蒙了一层灰纱。原因起始噪声 x_t 的采样分布和训练时不一致。我在采样代码里初始化x torch.randn_like(lq)但训练时noise torch.randn_like(gt)两者 shape 一致所以没报错可如果 lq 是经过 Resize 的它的分辨率跟 GT 不同噪声的频段分布就会有偏差。解决采样前强制把 lq resize 到和 GT 相同的尺寸再进网络。或者在数据加载阶段把 GT 和 lq 的尺寸设定为相同值不要 test 时单独用任意尺寸输入。U-Net 虽然能接受任意分辨率但扩散模型对输入分布极其敏感尺寸不匹配就是暗病。4.3 显存溢出batch size 调到 2 还是 OOM现象训练刚开始step 0 就报 CUDA out of memory直接中断。原因不是 U-Net 本身太占显存而是 attention 层在 256x256 分辨率下产生了巨大的中间特征。我的 base_channels128第四个 stage 是在 32x32 feature map 上做 attention这部分 memory 占用接近 8GB。解决把 attention 分辨率改成只在 16x16 层启用即 attention_resolution [16]会损失一点点高频细节但显存直降。另一个有效办法是开 gradient checkpointing用时间换显存代价是迭代变慢约 20%。4.4 恢复图产生重复纹理像原图被复印了好几份现象在真实低质图上测试时某些区域出现周期性重复结构尤其是纹理密集区域比如树丛、头发、草地。原因训练数据里低质图是由同一张 GT 加噪声合成的模型学会了“根据自身高频模式联想补充”也就是所谓的高频幻觉。真实退化图像的高频分布和合成噪声差异很大模型就开始了自由发挥。解决一是训练数据中混入多种退化类型不只高斯噪声加一部分下采样模糊、雨纹、JPEG 压缩噪声模型泛化性会明显提升二是推理时可以试试增加 DDIM steps 到 100更多步数意味着每一步恢复的幅度更小幻觉被逐步抑制。4.5 多卡训练 loss 降了但效果不如单卡现象用 DataParallel 跑双卡训练 loss 比单卡降得更快但最终采样效果无论是 PSNR 还是主观视觉都明显弱于单卡。原因BN 层在 DataParallel 里每个卡独立统计batch 变小导致统计量不稳定。扩散模型的 U-Net 里虽然没有显式 BN但 GroupNorm 在数据并行下有类似问题——不同卡的 normalization 统计不一致梯度更新互相干扰。解决直接用 DistributedDataParallel 替代 DataParallel每个进程独立维护完整的 batch 统计量。或者简单点放弃多卡单卡 3090 跑 300 epoch 大约 30 小时完全可接受。5. 实验流程与指标评估如何判断一组参数是否真的有效5.1 指标计算PSNR 和 SSIM 怎么算才是对的图像恢复领域最常规的客观指标是 PSNR 和 SSIM。但这里有个容易翻车的细节扩散模型输出的像素范围是 [-1,1]你和 GT 比较之前必须转回 [0,255] 的 uint8 空间如果直接在 float 张量上算 PSNR结果会整体偏高 24 dB这在论文里是绝对不允许的。这是我用的计算代码# utils/metrics.py def calculate_psnr_ssim(pred, gt, max_val255.0): pred pred.clamp(-1, 1) * 127.5 127.5 gt gt.clamp(-1, 1) * 127.5 127.5 pred pred.cpu().numpy().round().astype(np.uint8) gt gt.cpu().numpy().round().astype(np.uint8) psnr skimage.metrics.peak_signal_noise_ratio(gt, pred, data_rangemax_val) ssim skimage.metrics.structural_similarity( gt, pred, channel_axis-1, data_rangemax_val) return psnr, ssim注意structural_similarity的channel_axis参数在不同版本的 skimage 里这个位置参数有变化旧版是multichannelTrue新版是channel_axis-1。我一开始用旧参数在 scikit-image 0.22 上直接报错改成channel_axis-1就好了。PSNR 和 SSIM 不能只看绝对值要和你训练的 baseline 对比。比如你的输入 lq 本身 PSNR 是 22 dB恢复后到 28 dB这是有效提升如果只从 22 到 23.5那说明模型基本没学到东西只是在做轻度平滑。5.2 实验分组设计什么样的对比能说明模型真的有效我做扩散模型图像恢复实验时通常分成三组对比一组是输入 lq 直接双三次插值的结果作为最弱 baseline一组是用经典 U-Net 做回归预测的结果代表传统监督学习的上限一组是当前扩散模型的结果。同一份测试集上三组跑完再比较 PSNR / SSIM 和主观视觉。如果扩散模型的 PSNR 反而不如普通 U-Net 回归不要慌张。这是常见现象因为扩散模型的目标分布更接近真实高清图的流形而 PSNR 是逐像素误差GAN 和扩散模型这一类生成式方法在 PSNR 上通常拼不过直接回归方法但 SSIM 和主观效果更占优势。实际落地时我会同时报两组指标在论文或评审材料里写清楚扩散模型的优势在感知质量不在像素级误差。5.3 采样步数多少合适50 步与 250 步之间的权衡DDIM 采样步数直接决定推理耗时。我测试过同一 checkpoint 在不同 steps 下的表现10 步时图像有明显块状噪声30 步以上逐步改善50 步时和 100 步的 PSNR 差异在 0.1 dB 以内。考虑到 100 步推理耗时是 50 步的两倍实际部署我一般建议 50 步为上限对画质要求极高的场景再尝试 100 步。时间开销方面单张 3090 上50 步 DDIM 采样一张 256x256 图大约 2.5 秒50 步耗时约 1.2 秒。如果你的场景是离线批量处理这个速度可以接受如果是实时性要求高的场景得配合知识蒸馏或 latent diffusion 类方法进一步压缩步数那就超出这份代码的范围了。6. 进阶技巧用这个 DDPM 代码做超分和去雨的小改造这套代码最实用的点是它不需要伤筋动骨改动条件数据流就能迁移到超分和去雨任务。这里分享我实际改过的两条路径你拿到代码后可以照着切。6.1 从去噪切到超分只需要改数据配对和条件尺寸去噪任务的输入 lq 和输出 GT 天然同尺寸代码零改动。但超分任务里输入是下采样的小图输出是高清大图直接拼接会尺寸不匹配。我的做法是先把低分辨率图上采样到目标尺寸再拼接到噪声图上# 超分场景的输入处理 lq_resized F.interpolate(lq, size(256, 256), modebicubic, align_cornersFalse) x_cat torch.cat([x_t, lq_resized], dim1)这个做法本质上是让模型自己学会“如何补全上采样造成的模糊和缺失”。要注意modebicubic的选择不同插值方式影响模型接收到的条件信息质量。双三次是通用选择对 2x 超分效果好4x 超分时可以考虑modenearest保留边缘锐度模糊连带的信息反而更少。你可以在配置里调这个参数对比两组实验的 PSNR 差异。6.2 切到去雨任务控制合成雨纹的多样性去雨任务的退化模型是clean rain_streak雨纹的形状、方向、密度都影响模型的泛化。我合成训练数据时用了三条随机雨纹生成路径直线雨、弯曲雨、混合雨每张图随机选择 13 种叠加。重点在于不要让模型见到特定方向过硬的特征——雨纹方向最好在 [−60°, 60°] 之间随机旋转否则测试集换个方向雨纹就容易失效。训练超分和去雨时我一般把steps保持默认 1000eta0损失函数依旧 L1不需要额外改结构。唯一要调整的是数据集image_size如果你要处理 512x512 高清图显存不够就把 batch size 再降一档同时把随机 crop 改成从大图上采 256x256 的 patch 来训练推理时用滑窗拼接。这个思路在 video 领域叫 patch-based inference是验证大图效果最常用的手段。6.3 两分钟验证你的改法是否有效每次改完数据流程我都会做一个 5-epoch 的冒烟测试不跑完整训练。具体操作是在base.yaml里把epochs改成 5log_interval改成 10数据量限定在 200 张图训练完成后直接跑test.py看指标趋势python train.py --config configs/base.yaml --smoke_test python test.py --checkpoint ./checkpoints/model_step_2500.pth --test_dir ./data/test冒烟测试的关注点不是最终指标而是 loss 是否在下降、log 里的时间和显存是否符合预期以及采样输出的图像结构是否已经能看出内容轮廓。如果 5 epoch 后采样图还是一团噪声要么是学习率设置不合理要么是条件注入环节有 bug先别急着跑长训练把问题定位在源头。要说我自己的习惯从那以后我每次拿到一个新的扩散模型项目都强制自己先跑一遍冒烟测试再动完整训练。这个动作帮我挡掉了至少五六次因为数据集路径配错、归一化不一致导致的十几个小时白跑。希望这份代码和踩坑笔记也能帮你把第一个扩散模型图像恢复实验顺利跑通。本文还有配套的精品资源点击获取