简介这份资源面向深度学习研究者与图像恢复方向的开发者提供一套基于扩散模型diffusion model的完整可运行代码只需修改数据集路径即可直接用于去雨、去雾、去雪等多种图像恢复任务。压缩包共30个文件以13个Python源码为主辅以8个pyc编译文件、5个xml与1个yml配置、1个md说明等整体约29KB涵盖训练、采样、评估与模型定义等模块并附有Python版PSNR与SSIM计算脚本。资源内包含详细实验操作流程及参数路径修改方法训练与测试流程完整可跑代码中带有注释便于读者理解扩散模型在图像恢复中的实现细节也可直接迁移到自己的任务上训练测试。目前已有14986人学习下载适合希望快速上手扩散模型图像恢复、需要可复现实验代码的中高级读者参考使用。1. 扩散模型做图像恢复为什么它值得你花一个周末跑通图像恢复这件事传统做法是训练一个从退化图到清晰图的映射网络输入一张带噪、带雨、带雾或者低分辨率的图输出一张干净图。问题在于同一个退化输入往往对应多个合理输出L1/L2 损失会把模型逼向所有可能答案的平均值结果就是糊、平、缺纹理。扩散模型diffusion model换了个思路它不直接预测清晰图而是学习从纯噪声一步步去噪回清晰图的过程把「恢复」变成「条件生成」。你给它一张退化图当条件它就能在去噪的每一步里注入这张图的信息最终采样出一张既干净又有细节的结果。这套方案适合谁适合已经会写 PyTorch 训练循环、想把手里的图像恢复任务从「能跑」推到「好看」的人。它不适合完全没碰过深度学习的新手因为扩散模型的采样循环、噪声调度、条件注入方式都有坑调不通的时候你会怀疑人生。但只要你跑通过一次完整的训练加采样后面换数据集、换退化类型、换条件注入方式都是改几行配置的事。下面我按「先立住原理、再动手复现、最后讲坑」的顺序把一条能落地的路径讲清楚。2. 扩散模型用于图像恢复的原理与选型从 DDPM 到条件生成2.1 前向加噪与反向去噪到底在算什么扩散模型的核心是两个过程。前向过程是固定的对一张清晰图 $x_0$按预设的噪声调度逐步加高斯噪声经过 $T$ 步后变成近似纯噪声 $x_T$。这个过程可以写成闭式任意时刻 $t$ 的加噪结果 $x_t \sqrt{\bar\alpha_t} x_0 \sqrt{1-\bar\alpha_t}\epsilon$其中 $\epsilon$ 是标准高斯噪声$\bar\alpha_t$ 是累积的噪声系数。反向过程是学出来的训练一个网络 $\epsilon_\theta(x_t, t)$ 去预测第 $t$ 步加的噪声训练目标就是让预测噪声和真实噪声的均方误差最小。为什么这个目标能恢复图像因为网络学会了「给定一张带噪图噪声长什么样」采样时从纯噪声出发每一步减去网络预测的噪声逐步逼近清晰图。关键在于这个去噪过程是迭代的每一步都在修正所以它能生成纹理和细节而不是像单步回归那样输出平均值。2.2 图像恢复为什么必须做条件注入如果只训练无条件扩散模型你采样出来的是一张随机清晰图跟你的退化输入没关系。图像恢复要的是「以退化为条件」的生成所以必须把退化图 $y$ 作为条件喂给网络。常见做法有三种一是把 $y$ 和 $x_t$ 在通道维拼接后输入网络这是最直接的方式二是用单独的编码器提取 $y$ 的特征再通过交叉注意力注入去噪网络三是在采样阶段用退化模型做引导比如 DDRM 那类方法在频域做约束。我一般推荐第一种通道拼接改动最小训练最稳。它的代价是网络要同时处理噪声图和条件图显存占用翻倍但换来的是实现简单、调试直观。如果你做的是超分或者去雨这种退化类型固定的任务拼接方式足够用。2.3 噪声调度与采样步数的取舍DDPM 原始论文用 $T1000$ 步线性调度采样一次要跑 1000 次网络前向慢得让人抓狂。实际落地时大家普遍用余弦调度或者带偏移的线性调度并且把采样步数压到 50 到 200 步。DDIM 采样是另一个必选项它允许你跳步采样50 步就能出可接受的结果而且采样过程是确定性的方便复现。这里有个经验值图像恢复任务里采样步数低于 30 步时细节会明显丢失尤其是高频纹理50 到 100 步是质量和速度的平衡点超过 200 步收益很小纯属浪费算力。噪声调度方面余弦调度在低噪声阶段更平缓对恢复任务更友好我一般直接用它不折腾线性调度。2.4 网络骨架选 U-Net 还是 Transformer去噪网络的骨架选择目前主流还是 U-Net 加注意力机制。U-Net 的编码器-解码器结构天然适合多尺度特征跳跃连接能保留空间细节这对图像恢复很关键。Transformer 骨架比如 DiT在生成质量上有优势但训练成本高对数据量要求也大图像恢复这种任务用 U-Net 更划算。具体到 U-Net 的配置通道数从 64 起步每下采样一次翻倍通常做 3 到 4 次下采样。注意力层放在中间层和低分辨率层高分辨率层加注意力会爆显存。时间步 $t$ 的嵌入用正弦位置编码加两层 MLP注入到每个残差块里。这些配置在开源实现里都有成熟模板不用从零设计。3. 从零跑通图像恢复扩散模型环境、数据与训练脚本3.1 环境搭建与依赖版本锁定先把环境弄干净。扩散模型对 PyTorch 和 CUDA 版本比较敏感版本不匹配会出现各种玄学报错。我一般用 conda 建独立环境锁定以下版本组合这套组合在 30 系和 40 系卡上都跑通过。conda create -n diff_restore python3.10 -y conda activate diff_restore pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu118 pip install numpy1.24.3 pillow10.0.0 tqdm4.66.1 tensorboard2.14.0 pip install einops0.7.0逻辑说明torch 2.1.0 配 cu118 是当前比较稳的组合einops 用来做张量重排写注意力模块时比手写 permute 清晰。参数说明如果你用的是 40 系卡cu118 也兼容如果只有 CPU把 index-url 去掉装 CPU 版但训练会慢到没法用建议至少一张 8G 显存的卡。3.2 数据集准备与退化配对生成图像恢复需要成对数据清晰图 $x_0$ 和退化图 $y$。公开数据集里去噪用 BSD400 或 DIV2K去雨用 Rain100H超分用 DIV2K 的 bicubic 下采样版本。如果你有自己的数据按下面的脚本生成配对。import os import cv2 import numpy as np from glob import glob def make_pairs(clean_dir, save_dir, degrade_typegaussian, sigma25): os.makedirs(os.path.join(save_dir, clean), exist_okTrue) os.makedirs(os.path.join(save_dir, degraded), exist_okTrue) files glob(os.path.join(clean_dir, *.png)) glob(os.path.join(clean_dir, *.jpg)) for f in files: img cv2.imread(f) img cv2.resize(img, (256, 256)) name os.path.basename(f) if degrade_type gaussian: noise np.random.randn(*img.shape) * sigma degraded np.clip(img noise, 0, 255).astype(np.uint8) elif degrade_type blur: degraded cv2.GaussianBlur(img, (15, 15), 3) cv2.imwrite(os.path.join(save_dir, clean, name), img) cv2.imwrite(os.path.join(save_dir, degraded, name), degraded) make_pairs(./data/clean, ./data/pairs, degrade_typegaussian, sigma25)逻辑说明脚本把清晰图统一缩到 256×256按退化类型生成对应退化图分别存到 clean 和 degraded 两个文件夹。参数说明sigma 控制高斯噪声强度25 是去噪任务的常用值blur 分支的核大小 15 和标准差 3 是去模糊的常见配置。注意退化图要存成图片格式训练时再转张量这样方便你肉眼检查配对是否正确。3.3 条件扩散模型的核心训练循环训练循环是整个方案的心脏。下面这段代码实现了通道拼接的条件注入、余弦噪声调度和简化的损失计算。import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange class ConditionalUNet(nn.Module): def __init__(self, in_ch6, base_ch64): super().__init__() # 输入通道是63通道噪声图 3通道条件图 self.enc1 nn.Conv2d(in_ch, base_ch, 3, padding1) self.enc2 nn.Conv2d(base_ch, base_ch * 2, 3, stride2, padding1) self.mid nn.Conv2d(base_ch * 2, base_ch * 2, 3, padding1) self.dec1 nn.ConvTranspose2d(base_ch * 2, base_ch, 2, stride2) self.dec2 nn.Conv2d(base_ch, 3, 3, padding1) self.time_mlp nn.Sequential(nn.Linear(128, base_ch * 2), nn.SiLU()) def forward(self, x_t, cond, t): # x_t: 噪声图, cond: 退化条件图, t: 时间步 h torch.cat([x_t, cond], dim1) t_emb self.time_mlp(self.pos_emb(t)) h1 F.silu(self.enc1(h)) h2 F.silu(self.enc2(h1)) t_emb[:, :, None, None] h3 F.silu(self.mid(h2)) d1 F.silu(self.dec1(h3)) out self.dec2(d1 h1) return out def pos_emb(self, t): # 正弦位置编码把标量时间步变成128维向量 half 64 freqs torch.exp(torch.linspace(0, 9, half, devicet.device)) args t[:, None].float() * freqs[None] return torch.cat([torch.sin(args), torch.cos(args)], dim-1) def cosine_alpha_bar(T1000): # 余弦噪声调度返回累积alpha steps torch.linspace(0, T, T 1) f torch.cos((steps / T 0.008) / 1.008 * torch.pi / 2) ** 2 return f / f[0] def train_step(model, x0, cond, optimizer, alpha_bar): b x0.shape[0] t torch.randint(0, len(alpha_bar) - 1, (b,), devicex0.device) ab alpha_bar[t][:, None, None, None] eps torch.randn_like(x0) x_t torch.sqrt(ab) * x0 torch.sqrt(1 - ab) * eps pred model(x_t, cond, t) loss F.mse_loss(pred, eps) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item()逻辑说明网络输入是噪声图和条件图的通道拼接时间步通过正弦编码加 MLP 注入到中间层。训练时随机采样时间步 $t$按闭式公式构造 $x_t$让网络预测噪声损失是预测噪声和真实噪声的 MSE。参数说明base_ch 控制模型宽度64 是入门配置显存够可以上 128T1000 是训练时的总步数采样时可以跳步余弦调度的 0.008 是防止 $t0$ 附近噪声过小的偏移量别改。3.4 采样脚本与 DDIM 跳步实现训练完之后采样是验证效果的关键。DDIM 采样允许你用远少于训练步数的迭代出图。torch.no_grad() def ddim_sample(model, cond, alpha_bar, steps50, eta0.0): # eta0 是确定性采样eta1 退化成DDPM device cond.device x torch.randn_like(cond) ts torch.linspace(len(alpha_bar) - 1, 0, steps 1).long() for i in range(steps): t_cur, t_next ts[i], ts[i 1] ab_cur alpha_bar[t_cur] ab_next alpha_bar[t_next] if t_next 0 else torch.tensor(1.0, devicedevice) t_batch torch.full((cond.shape[0],), t_cur, devicedevice, dtypetorch.long) eps_pred model(x, cond, t_batch) x0_pred (x - torch.sqrt(1 - ab_cur) * eps_pred) / torch.sqrt(ab_cur) x0_pred x0_pred.clamp(-1, 1) sigma eta * torch.sqrt((1 - ab_next) / (1 - ab_cur) * (1 - ab_cur / ab_next)) noise torch.randn_like(x) if eta 0 else 0 x torch.sqrt(ab_next) * x0_pred torch.sqrt(1 - ab_next - sigma ** 2) * eps_pred sigma * noise return x逻辑说明DDIM 的核心是先由当前 $x_t$ 和预测噪声反推 $x_0$再用 $x_0$ 和预测噪声重构下一步的 $x_{t-1}$。eta 控制随机性图像恢复任务一般设 0保证输出稳定可复现。参数说明steps 设 50 是速度和质量的平衡点clamp 到 [-1,1] 是因为训练时图像归一化到了这个范围不 clamp 会出现色彩溢出。4. 训练与采样中的避坑排查五个血泪教训4.1 损失不降反升采样全是噪声现象训练前几百步 loss 正常下降之后开始震荡上升采样出来是一团噪声。原因学习率设太大或者噪声调度的 $\bar\alpha_t$ 计算有误导致 $x_t$ 的噪声比例不对。解决先把学习率降到 1e-4 甚至 5e-5用余弦调度重新算 alpha_bar打印几个时间步的 $\bar\alpha_t$ 确认单调递减。我踩过这个坑当时是 alpha_bar 的索引写反了调了两天才发现。4.2 采样结果和输入退化图完全无关现象不管输入什么退化图采样结果都差不多像是无条件生成。原因条件图没有正确拼接到网络输入或者条件图在预处理时被归一化到了错误范围。解决检查 forward 里的 torch.cat 维度是不是通道维确认条件图和噪声图用了相同的归一化参数。常见错误是清晰图归一化到 [-1,1]退化图却留在 [0,255]网络根本学不到对应关系。4.3 显存爆炸batch size 上不去现象batch size 设 4 就 OOM但模型参数量看起来不大。原因扩散模型训练要同时存 $x_0$、$x_t$、条件图、预测噪声和中间激活显存占用是普通回归网络的三倍左右。解决用混合精度训练把 batch size 降到 2 或 1配合梯度累积。另外注意力层放在高分辨率特征图上会急剧增加显存把注意力限制在 32×32 以下的分辨率。4.4 采样步数减少后细节糊成一片现象用 DDIM 把步数从 1000 降到 20速度上去了但纹理全没了。原因跳步太大去噪过程在低噪声阶段跨过了太多细节恢复的步骤。解决步数不要低于 50或者在低噪声阶段用更密的步长。我一般用非均匀步长前 80% 的步数覆盖高噪声阶段后 20% 覆盖低噪声阶段细节保留明显更好。4.5 训练集上效果好换一张图就翻车现象在训练集图片上恢复效果很好拿一张没见过的图测试结果出现明显伪影。原因过拟合或者退化类型不匹配。解决数据增强必须做随机裁剪、翻转、旋转都要上退化类型要覆盖测试时的真实退化如果测试图是真实噪声而不是高斯噪声训练时就得混入真实噪声样本。别指望模型能泛化到没见过的退化分布。5. 把扩散图像恢复推到可用进阶技巧与验证方法5.1 用感知损失和对抗损失补细节纯 MSE 训练出来的模型PSNR 可能不错但肉眼看着就是糊。原因是 MSE 对高频细节不敏感。进阶做法是在训练后期加入感知损失用预训练 VGG 提取特征算 L1 距离再加一个轻量判别器做对抗训练。感知损失让输出在特征空间接近清晰图对抗损失逼出真实纹理。注意对抗损失权重不能大一般设 0.01 到 0.05大了会引入伪影。5.2 验证不能只看 PSNRPSNR 和 SSIM 是图像恢复的常规指标但它们和肉眼观感经常不一致。我一般同时看三个东西PSNR/SSIM 数值、LPIPS 感知距离、以及局部放大对比图。LPIPS 越低说明感知质量越好通常和肉眼判断更一致。验证集要固定一批图每次改配置都跑同一批不然指标波动你分不清是模型变了还是数据变了。5.3 一个具体技巧从预训练无条件模型微调从零训练条件扩散模型很费卡。更省事的路径是拿一个预训练的无条件扩散模型把条件图通过 ControlNet 那种旁路注入只训练旁路分支。这样训练量减少一个数量级效果还更稳。具体做法是冻结原模型权重复制一份编码器作为条件编码器在中间层用零卷积把条件特征加到原模型上。零卷积初始输出为零保证微调初期不破坏原模型能力。验证项工具合格线参考PSNR自写脚本去噪任务 30dB 以上SSIMskimage0.85 以上LPIPSlpips 库0.15 以下采样耗时time 模块50 步单张 2 秒内这套方案我从头跑通过三次换过去噪、去雨、超分三种任务最深的教训是别一上来就堆大模型先把 64 通道的小模型在 256×256 上跑通确认训练循环、噪声调度、条件注入都没问题再往上加。扩散模型的坑大多不在模型结构而在噪声调度和条件对齐这些细节上。希望帮到你。本文还有配套的精品资源点击获取