简介面向深度学习图像恢复需求提供一套基于扩散模型diffusion model的完整可运行代码适用于去雨、去雾、去雪等多类降质场景适合具备一定PyTorch基础的研究者直接迁移到自定义数据集。压缩包共30个文件以Python源码为主含模型结构、训练/测试脚本、指标计算等13个py文件另有配置文件、环境配置及说明文档整体仅29KB轻量易部署。已有14979人学习此资源。代码只需修改数据集路径即可在不同图像恢复任务间切换训练与测试流程完整同时附带详细实验操作流程和参数路径修改方法内置常用的PSNR、SSIM计算脚本关键位置提供注释并结合配套博客辅助理解可帮助读者快速上手并复现实验。1. 为什么图像恢复要选扩散模型从“一步映射”到“迭代修正”做图像去噪、去模糊或超分时很多人一开始会把“扩散模型diffusion model用于图像恢复”理解成大一号的神经网络回归给一张退化图输出一张干净图。真正把完整可运行代码跑一遍就会发现它和你熟悉的SRGAN、Restormer完全不是同一个思路——扩散模型不学“从退化图到干净图”的直接映射而是学“把图一步步弄脏再学习反着弄干净”的过程。它能解决高斯噪声、模糊、低分辨率等恢复问题尤其适合追求纹理细节和分布真实的场景。下面的落地路径适合有PyTorch基础、想在自己数据集上跑通并继续改进扩散恢复项目的工程师我按“怎么想→怎么写→怎么跑→怎么排坑→怎么验证”来讲每一步都是可执行的。2. 从DDPM到条件扩散先弄懂模型在恢复任务里到底学什么2.1 前向加噪与反向去噪扩散模型的最小闭环扩散模型的起点是DDPM里的两个过程前向过程按固定噪声表把干净图x0逐步变成纯噪声反向过程让网络学会每一步“去掉多少噪声”。做图像恢复时不改变这个闭环改变的是网络输入上多挂一个退化图y。先看前向。前向公式可以写成x_t sqrt(alpha_bar_t) * x0 sqrt(1 - alpha_bar_t) * epsilonepsilon是标准高斯噪声。实际编码时不要循环t步去模拟加噪而是每次都从x0直接“跳”到第t步因为多步叠加的最终结果有解析表达式。下面这个函数就是完整加噪逻辑import numpy as np def forward_diffusion_sample(x0, t, alphas_cumprod): 把干净图 x0 加噪到第 t 步。 x0: 形状 (C, H, W)建议值域在 [-1, 1] t: 整数0 t len(alphas_cumprod) sqrt_alpha_bar np.sqrt(alphas_cumprod[t]) sqrt_one_minus np.sqrt(1.0 - alphas_cumprod[t]) noise np.random.randn(*x0.shape).astype(np.float32) x_t sqrt_alpha_bar * x0 sqrt_one_minus * noise return x_t, noise函数返回x_t和这次加的noise。训练时网络的预测目标就是noise而不是直接预测x0。为什么用alphas_cumprod数组因为DDPM加噪过程是马尔可夫链t步的累积效果可以解析表达训练时不必仿真t步直接跳到任意时刻即可能省下大量计算。alphas_cumprod通常由beta schedule换算而来beta从1e-4线性涨到0.02t越大alpha_bar越接近0x_t越接近纯噪声t0时x_t就是x0。如果你的数据值域在[0,1]而不是[-1,1]这套公式的方差匹配会出问题后面避坑章会单独讲。反向去噪就是逆向迭代从x_T纯噪声出发每一步用网络估出的噪声“去噪”一步逐步还原x0。去掉噪声的细节公式放在第4章的采样代码里那里能看到完整流程。这里先记住一件事前向加噪的反过程不是简单解卷积而是由网络一步步估计条件分布这也是扩散模型能生成清晰高频纹理的根本原因。2.2 从无条件到条件为什么图像恢复要用条件扩散无条件扩散学的是数据分布p(x)也就是“生成一张看起来真实的图”。图像恢复的目标不同给定退化观测y要恢复出对应的x0应该建模p(x0|y)。工程实现上把退化图y作为条件输入网络最常见的做法是通道拼接网络输入变成concat(x_t, y)在通道维上把当前含噪图和退化图并列。为什么这样能恢复因为x_t携带的是“全局结构当前噪声水平”y携带的是“真实观测约束和细节线索”。训练时让网络同时看到这两类信息来源学出来的去噪步会趋向于“在保持y语义的前提下去除噪声”而不是把图重构到任意一个真实样本。这也是为什么训练和推理时y必须来自同样的退化管线——训练时用sigma25的高斯噪声推理时拿sigma50的图来跑模型会不知所措。条件注入除了concat还有把条件编码后加到timestep embedding上以及cross-attention方案。对图像恢复这种像素级任务我一般建议先用channel-wise concat代码最简单、也最容易排查网络输入通道数等于图像通道数乘以2。等这个跑通了再对比其他注入方式。直接上复杂条件结构容易把问题引入黑匣子肉眼根本分不清是生成问题还是条件注入问题。判断模型写没写对最好的办法就是先构造随机输入跑一次forward别急着训。2.3 选型理由为什么图像恢复常用DDPM/Gaussian diffusion而不是别的选DDPM做图像恢复有三个现实理由。第一是退化类型恰好对齐图像恢复里最核心的高斯去噪其前向过程本身就是高斯加噪扩散模型天然匹配这不是比喻是数学上的同构。第二是质量与调参的平衡和GAN相比扩散模型训练更稳不容易模式坍缩loss曲线一眼能看明白对新手友好和纯回归网络Restormer、SwinIR这类相比扩散模型在纹理修复上通常更真实代价是多步采样耗时大。第三是工程起点低一个能跑的DDPM图像恢复核心代码可以控制在300行上下数据准备、训练、采样逻辑都简单适合先验证自己的idea确认方向有效再引入LDM、ControlNet等变体。这里要特别提醒不要一上来就想上stable diffusion那套latent diffusion结构。图像恢复里大多数实验仍直接在像素空间训练DDPM原因是恢复任务目标是恢复图像本身不希望latent重建过程丢细节像素空间diffusion代码少、更容易加条件。只有当训练分辨率超过256×256、或数据量很大时才考虑迁移到LDM。另外如果你习惯先看示例代码讲解建议按“模型前向→训练循环→采样”的顺序读代码而不是从头到尾读文件——这样能最快把各文件串起来。3. 把完整可运行代码跑起来文件结构、依赖与最小训练脚本3.1 项目文件结构与依赖我习惯把一次扩散恢复实验拆成五个文件保证换数据集、换退化类型时不用重写模型代码diffusion_restore/ ├── config.yaml # 噪声表、步数、学习率、数据路径等 ├── dataset.py # 读取干净图/退化图返回成对数据 ├── model.py # U-Net timestep embedding ├── train.py # 训练主循环含EMA └── sample.py # 从退化图开始恢复的采样脚本config.yaml把超参独立出来后面做消融实验时只改配置不碰代码这个习惯能省一半时间。依赖方面我一般固定这几样Python 3.8以上、PyTorch 2.1以上、NumPy、OpenCV-Python、tqdm。训练一张常见显卡如RTX 3060 12G就能跑128×128的batch 8不需要多卡。3.2 用PyTorch实现U-Net骨架的关键部分model.py里最核心的是两块timestep embedding和条件输入。下面给一个能跑通最小规模的骨架真正实验时把down/up部分替换成你惯用的U-Net结构即可import torch import torch.nn as nn import math class TimestepEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim dim self.mlp nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim) ) def forward(self, t): half self.dim // 2 emb math.log(10000) / (half - 1) emb torch.exp(torch.arange(half, devicet.device) * -emb) emb t[:, None] * emb[None, :] emb torch.cat([torch.sin(emb), torch.cos(emb)], dim1) return self.mlp(emb) class SimpleRestoreUNet(nn.Module): def __init__(self, in_channels6, base_dim64): super().__init__() self.in_conv nn.Conv2d(in_channels, base_dim, 3, padding1) self.mid_conv nn.Conv2d(base_dim, base_dim * 2, 3, padding1) self.out_conv nn.Conv2d(base_dim * 2, 3, 3, padding1) self.t_emb TimestepEmbedding(base_dim) self.gamma nn.Linear(base_dim, base_dim) def forward(self, x_t, y, t): x torch.cat([x_t, y], dim1) # 条件拼接 t_emb self.t_emb(t) h self.in_conv(x) w self.gamma(t_emb).view(-1, h.shape[1], 1, 1) h h * (1 w) # 时间步信息注入 h torch.relu(self.mid_conv(h)) return self.out_conv(h)这段代码有几个关键点。第一in_channels6是因为彩色图的x_t和退化图y各3个通道拼接后是6灰度恢复任务则改成2。第二网络输出通道仍是3因为预测目标是与输入图同尺寸的噪声图。第三这里用了一个简化U-Net骨架真正实验时要替换为带4次下采样/上采样和skip connection的结构时间步注入也可以改成scaleshift的AdaGN形式。判断模型是否写对可以用model(torch.randn(2,3,64,64), torch.randn(2,3,64,64), torch.randint(0,1000,(2,)))跑一次forward输出形状应为(2,3,64,64)。3.3 训练主循环loss计算、EMA、学习率、日志train.py主循环是完整可运行代码的核心。注意两个细节加噪的干净图来自clean目标退化图只作为条件loss是预测噪声和真实噪声的MSE。for step, (clean, degraded) in enumerate(dataloader): clean clean.to(device) degraded degraded.to(device) t torch.randint(0, num_timesteps, (clean.size(0),), devicedevice) x_t, noise forward_diffusion_sample(clean, t, alphas_cumprod) pred_noise model(x_t, degraded, t) loss torch.nn.functional.mse_loss(pred_noise, noise) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema_update(model, ema_model, decay0.999) if step % 50 0: print(fstep {step:6d} | loss {loss.item():.6f})训练时用x_t从clean加噪而不是从degraded加噪原因是目标噪声必须和输入严格对应模型学“把加噪后的clean还原”的同时参考degraded做条件约束。EMA维护一组更平滑的权重推理时用ema_model通常比model高0.2~0.5dB值得一开始就加上。clip_grad_norm_防止早期loss爆炸学习率一般用1e-4AdamW配合cosine schedule。把dataset.py、model.py、train.py、sample.py按上面的关键函数填入就是一份能跑通128×128高斯去噪的最小完整代码。4. 详细实验操作流程从数据准备到完整采样实验4.1 数据准备与退化模拟去噪为例实验流程第一步是造“成对数据”。用公开的BSD400或DIV2K数据集预先切成128×128的patch退化图现场加高斯噪声sigma25/255是常用配置。加噪与裁剪的示例代码import cv2 import numpy as np def add_gaussian_noise(img, sigma25): # img: float32值域 [0,1]sigma 以像素值为单位 noise np.random.randn(*img.shape).astype(np.float32) * (sigma / 255.0) return np.clip(img.astype(np.float32) noise, 0.0, 1.0) def random_crop_pair(clean, degraded, crop_size128): H, W clean.shape[:2] top np.random.randint(0, H - crop_size 1) left np.random.randint(0, W - crop_size 1) return (clean[top:topcrop_size, left:leftcrop_size], degraded[top:topcrop_size, left:leftcrop_size])sigma这里有个最容易踩的暗坑如果输入值域是[-1,1]加噪标准差应为sigma/127.5这是把噪声幅度相对值域做映射。很多第一次复现的人在这里翻车值域用[0,1]sigma却直接写25最后采样结果整体发灰。所以我建议dataset.py里统一走[-1,1]读图后执行img img / 127.5 - 1这和DDPM加噪公式天然一致。完整的数据集类按PyTorch惯例写读入图片、在线加噪、返回clean和degraded即可。退化模拟必须和你的恢复目标一致做去噪就加高斯噪声做超分就下采样再上采样做去模糊就要模糊核。训练时用什么退化参数测试时就用什么参数这是实验可复现的基础。4.2 训练配置与超参数一张表看懂最常用的起步值起步阶段的参数我一般直接照下表填参数推荐值说明num_timesteps1000DDPM标准设置显存紧张可降到500beta_schedulelinear 1e-4→0.02最常用噪声表高分辨率可换cosinecrop_size128兼顾纹理表达与显存占用batch_size812G显存可跑16G可提到16learning_rate1e-4AdamWcosine衰减ema_decay0.999推理权重用EMA版本grad_clip1.0防前期loss爆炸epochs100数据量小时建议数step而不是epoch这些值不是拍脑袋定的num_timesteps降到500时alphas_cumprod要重新生成采样步数也相应降到100左右crop_size不建议小于64否则纹理信息不够去噪结果容易糊batch_size和显存关系密切不够就减crop或开AMP。4.3 推理采样流程从纯噪声到恢复图的确定性采样训练完成后sample.py按“从噪声出发迭代去噪”的方式做恢复。下面这段是DDIM式更新用num_steps控制实际采样步数def sample_restore(model, degraded, num_steps200, eta0.0): model.eval() x torch.randn_like(degraded) # 从纯噪声图开始 with torch.no_grad(): for i in reversed(range(num_steps)): t i * (num_timesteps // num_steps) # 映射到训练步 t_batch torch.full((degraded.size(0),), t, devicedegraded.device) pred_noise model(x, degraded, t_batch) alpha_bar alphas_cumprod[t] alpha_bar_prev alphas_cumprod[t-1] if t 0 else 1.0 x0_pred (x - torch.sqrt(1 - alpha_bar) * pred_noise) / torch.sqrt(alpha_bar) x0_pred x0_pred.clamp(-1, 1) sigma_t eta * torch.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * torch.sqrt(1 - alpha_bar / alpha_bar_prev) x torch.sqrt(alpha_bar_prev) * x0_pred \ torch.sqrt(1 - alpha_bar_prev - sigma_t**2) * pred_noise \ sigma_t * torch.randn_like(x) return x0_pred这段代码用eta0的确定性更新每一步先用pred_noise反推出“干净图估计”x0_pred并clamp再按前向公式把这个估计重新加噪到下一步循环结束后返回最后的x0_pred。num_steps200是速度和质量比较稳的平衡点常见实验用100~200步想先验证pipeline可以跑100步。eta这个参数很关键eta0时没有随机噪声采样是确定性的多次推理结果一致方便调试eta0引入随机性可能带来更丰富的纹理但也会让结果不稳定。推理输入degraded必须做和训练完全一致的值域转换否则x0_pred第一步就会被clamp拉到错误区间。提示采样输出记得从[-1,1]映射回[0,1]再保存((x0_pred 1) / 2).clamp(0, 1)。这一步漏了后面算PSNR/SSIM全都会错位。5. 避坑与常见问题排查训练不收敛、图像发灰、采样发糊的五个现场5.1 loss在降采样图却“脏”先查条件拼接有没有配对错现象训练loss稳步下降采样结果却像“噪点色块”的混合体完全不是清晰的恢复图。原因最常见的是训练循环里把x_t误取自degraded而不是clean。目标噪声是“对clean加噪得到的噪声”输入x_t却是“对degraded加噪”两者语义不一致网络学的不是去噪而是噪声混淆。这种情况你用代码诊断插件去抓往往第一个暴露的就是数据流不对。解决把train循环里的forward_diffusion_sample(clean, t, alphas_cumprod)打出来核对再把model(x_t, degraded, t)改成model(x_t, torch.zeros_like(degraded), t)跑一次debug如果去掉条件还能学到基本去噪说明网络没问题问题一定在条件拼接。5.2 图像发灰、对比度低检查值域与clamp现象恢复图看起来“雾蒙蒙”细节不够锐利肉眼明显发灰但PSNR不算太低。原因大概率是图像值域和网络假设不一致。模型内部按[-1,1]设计代码里却把[0,1]的图直接喂进去或者采样时x0_pred被clamp到[-1,1]但前处理没有把图像归一到这个区间导致每一步的估计都系统性偏灰。解决固定统一值域。读图后立即img img / 127.5 - 1采样结束后out (out 1) / 2加噪噪声标准差按sigma/127.5计算。这个坑在换数据集时最容易复发建议直接写进dataset.py的模块注释里。5.3 训练到一半显存溢出或速度极慢现象batch16、crop128跑几十个step就OOM或者一晚上只跑完几千step明显不对劲。原因diffusion模型每步要同时维护x_t、noise、y和梯度显存占用比普通CNN高不少另外可能是没开混合精度cudnn benchmark也没开卷积反复autotune浪费大量时间。解决优先降crop到96或64batch减半然后开AMP用torch.autocast(device_typecuda, dtypetorch.float16)包forward/backward配合torch.amp.GradScaler。显存实在小就把una的base_dim从64降到32。注意力模块如果实现不当也会吃显存但用第3章那个简化骨架不会有这个问题。5.4 采样步数200步太慢减少后效果明显劣化现象200步出一张128×128图要一分多钟降成50步后图变糊或出现噪点。原因两个可能。一是t映射没对齐直接把reversed(range(num_steps))当成t输入网络而不是按比例映射到1000步空间二是更新公式里x0_pred没有正确clamp步数减少后误差累积。解决先确认映射是t i * (num_timesteps // num_steps)。然后对比100、150、200步的PSNR找到质量拐点。若必须快速采样再考虑DDIM、DPM-Solver或consistency蒸馏方向——那是后话别指望靠改一两个参数既快又好。5.5 训练集指标高、测试集崩盘数据量不够时的过拟合现象训练集上PSNR已经32dB以上测试集只有24~25dB两者差距巨大。原因扩散模型本质是生成模型数据量太小时会直接背训练集分布小数据集配上大batch这个问题会暴露得更快。解决用128×128随机裁剪把每张图变成几十个patch训练集用BSD400DIV2K混合并行加入随机翻转和旋转。如果条件允许先在无监督大规模噪声集合上预训练再在下游数据上微调。注意千万不要把测试集的裁剪图混进训练这是最常见的评测污染。6. 用固定测试集、固定随机种子做实验三个把结果做扎实的验证习惯先看固定随机种子的代码这是所有对比实验的地基import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False第一个习惯每次训练结束必须用同一个测试子集、同一个seed重新采样。这样不同配置之间的对比才是公平的。如果每次推理都放任随机源变化哪怕同一份权重两次采样结果也会有肉眼可见的差异指标对比毫无意义。第二个习惯评估时统一用skimage.metrics.peak_signal_noise_ratio和structural_similarity但一定要小心data_range参数。图像值域[-1,1]时data_range2.0值域[0,1]时data_range1.0。我第一次复现去噪实验时PSNR比baseline低0.5dB查了一天发现是data_range写成了1.0而模型输出是[-1,1]。这类“低分”往往不是模型问题是评测口径问题。第三个习惯每次实验记录轻量日志。我在每次训练前会把config.yaml内容、数据集版本、退化参数、最终PSNR/SSIM这四样复制进一个CSV同一配置至少跑两次保留seed记录。扩散模型采样有随机性一次性结果不可信。每次多花10分钟记录能少走一周弯路。另外可以固定一个“小验证集”训练每500步采样3张图肉眼观察恢复质量。loss曲线只能告诉你是否收敛不能告诉你恢复效果是否真实在屏幕上花20秒看的直观反馈比任何指标都早。我自己后来养成的习惯是任何新配置先不跑完整训练先在测试集上拿20步采样跑通一次全链路再放开训练等链路通了再加速调参。希望这个习惯也能帮到你。本文还有配套的精品资源点击获取 扩散模型图像恢复实战：从DDPM原理到PyTorch代码
简介面向深度学习图像恢复需求提供一套基于扩散模型diffusion model的完整可运行代码适用于去雨、去雾、去雪等多类降质场景适合具备一定PyTorch基础的研究者直接迁移到自定义数据集。压缩包共30个文件以Python源码为主含模型结构、训练/测试脚本、指标计算等13个py文件另有配置文件、环境配置及说明文档整体仅29KB轻量易部署。已有14979人学习此资源。代码只需修改数据集路径即可在不同图像恢复任务间切换训练与测试流程完整同时附带详细实验操作流程和参数路径修改方法内置常用的PSNR、SSIM计算脚本关键位置提供注释并结合配套博客辅助理解可帮助读者快速上手并复现实验。1. 为什么图像恢复要选扩散模型从“一步映射”到“迭代修正”做图像去噪、去模糊或超分时很多人一开始会把“扩散模型diffusion model用于图像恢复”理解成大一号的神经网络回归给一张退化图输出一张干净图。真正把完整可运行代码跑一遍就会发现它和你熟悉的SRGAN、Restormer完全不是同一个思路——扩散模型不学“从退化图到干净图”的直接映射而是学“把图一步步弄脏再学习反着弄干净”的过程。它能解决高斯噪声、模糊、低分辨率等恢复问题尤其适合追求纹理细节和分布真实的场景。下面的落地路径适合有PyTorch基础、想在自己数据集上跑通并继续改进扩散恢复项目的工程师我按“怎么想→怎么写→怎么跑→怎么排坑→怎么验证”来讲每一步都是可执行的。2. 从DDPM到条件扩散先弄懂模型在恢复任务里到底学什么2.1 前向加噪与反向去噪扩散模型的最小闭环扩散模型的起点是DDPM里的两个过程前向过程按固定噪声表把干净图x0逐步变成纯噪声反向过程让网络学会每一步“去掉多少噪声”。做图像恢复时不改变这个闭环改变的是网络输入上多挂一个退化图y。先看前向。前向公式可以写成x_t sqrt(alpha_bar_t) * x0 sqrt(1 - alpha_bar_t) * epsilonepsilon是标准高斯噪声。实际编码时不要循环t步去模拟加噪而是每次都从x0直接“跳”到第t步因为多步叠加的最终结果有解析表达式。下面这个函数就是完整加噪逻辑import numpy as np def forward_diffusion_sample(x0, t, alphas_cumprod): 把干净图 x0 加噪到第 t 步。 x0: 形状 (C, H, W)建议值域在 [-1, 1] t: 整数0 t len(alphas_cumprod) sqrt_alpha_bar np.sqrt(alphas_cumprod[t]) sqrt_one_minus np.sqrt(1.0 - alphas_cumprod[t]) noise np.random.randn(*x0.shape).astype(np.float32) x_t sqrt_alpha_bar * x0 sqrt_one_minus * noise return x_t, noise函数返回x_t和这次加的noise。训练时网络的预测目标就是noise而不是直接预测x0。为什么用alphas_cumprod数组因为DDPM加噪过程是马尔可夫链t步的累积效果可以解析表达训练时不必仿真t步直接跳到任意时刻即可能省下大量计算。alphas_cumprod通常由beta schedule换算而来beta从1e-4线性涨到0.02t越大alpha_bar越接近0x_t越接近纯噪声t0时x_t就是x0。如果你的数据值域在[0,1]而不是[-1,1]这套公式的方差匹配会出问题后面避坑章会单独讲。反向去噪就是逆向迭代从x_T纯噪声出发每一步用网络估出的噪声“去噪”一步逐步还原x0。去掉噪声的细节公式放在第4章的采样代码里那里能看到完整流程。这里先记住一件事前向加噪的反过程不是简单解卷积而是由网络一步步估计条件分布这也是扩散模型能生成清晰高频纹理的根本原因。2.2 从无条件到条件为什么图像恢复要用条件扩散无条件扩散学的是数据分布p(x)也就是“生成一张看起来真实的图”。图像恢复的目标不同给定退化观测y要恢复出对应的x0应该建模p(x0|y)。工程实现上把退化图y作为条件输入网络最常见的做法是通道拼接网络输入变成concat(x_t, y)在通道维上把当前含噪图和退化图并列。为什么这样能恢复因为x_t携带的是“全局结构当前噪声水平”y携带的是“真实观测约束和细节线索”。训练时让网络同时看到这两类信息来源学出来的去噪步会趋向于“在保持y语义的前提下去除噪声”而不是把图重构到任意一个真实样本。这也是为什么训练和推理时y必须来自同样的退化管线——训练时用sigma25的高斯噪声推理时拿sigma50的图来跑模型会不知所措。条件注入除了concat还有把条件编码后加到timestep embedding上以及cross-attention方案。对图像恢复这种像素级任务我一般建议先用channel-wise concat代码最简单、也最容易排查网络输入通道数等于图像通道数乘以2。等这个跑通了再对比其他注入方式。直接上复杂条件结构容易把问题引入黑匣子肉眼根本分不清是生成问题还是条件注入问题。判断模型写没写对最好的办法就是先构造随机输入跑一次forward别急着训。2.3 选型理由为什么图像恢复常用DDPM/Gaussian diffusion而不是别的选DDPM做图像恢复有三个现实理由。第一是退化类型恰好对齐图像恢复里最核心的高斯去噪其前向过程本身就是高斯加噪扩散模型天然匹配这不是比喻是数学上的同构。第二是质量与调参的平衡和GAN相比扩散模型训练更稳不容易模式坍缩loss曲线一眼能看明白对新手友好和纯回归网络Restormer、SwinIR这类相比扩散模型在纹理修复上通常更真实代价是多步采样耗时大。第三是工程起点低一个能跑的DDPM图像恢复核心代码可以控制在300行上下数据准备、训练、采样逻辑都简单适合先验证自己的idea确认方向有效再引入LDM、ControlNet等变体。这里要特别提醒不要一上来就想上stable diffusion那套latent diffusion结构。图像恢复里大多数实验仍直接在像素空间训练DDPM原因是恢复任务目标是恢复图像本身不希望latent重建过程丢细节像素空间diffusion代码少、更容易加条件。只有当训练分辨率超过256×256、或数据量很大时才考虑迁移到LDM。另外如果你习惯先看示例代码讲解建议按“模型前向→训练循环→采样”的顺序读代码而不是从头到尾读文件——这样能最快把各文件串起来。3. 把完整可运行代码跑起来文件结构、依赖与最小训练脚本3.1 项目文件结构与依赖我习惯把一次扩散恢复实验拆成五个文件保证换数据集、换退化类型时不用重写模型代码diffusion_restore/ ├── config.yaml # 噪声表、步数、学习率、数据路径等 ├── dataset.py # 读取干净图/退化图返回成对数据 ├── model.py # U-Net timestep embedding ├── train.py # 训练主循环含EMA └── sample.py # 从退化图开始恢复的采样脚本config.yaml把超参独立出来后面做消融实验时只改配置不碰代码这个习惯能省一半时间。依赖方面我一般固定这几样Python 3.8以上、PyTorch 2.1以上、NumPy、OpenCV-Python、tqdm。训练一张常见显卡如RTX 3060 12G就能跑128×128的batch 8不需要多卡。3.2 用PyTorch实现U-Net骨架的关键部分model.py里最核心的是两块timestep embedding和条件输入。下面给一个能跑通最小规模的骨架真正实验时把down/up部分替换成你惯用的U-Net结构即可import torch import torch.nn as nn import math class TimestepEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim dim self.mlp nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim) ) def forward(self, t): half self.dim // 2 emb math.log(10000) / (half - 1) emb torch.exp(torch.arange(half, devicet.device) * -emb) emb t[:, None] * emb[None, :] emb torch.cat([torch.sin(emb), torch.cos(emb)], dim1) return self.mlp(emb) class SimpleRestoreUNet(nn.Module): def __init__(self, in_channels6, base_dim64): super().__init__() self.in_conv nn.Conv2d(in_channels, base_dim, 3, padding1) self.mid_conv nn.Conv2d(base_dim, base_dim * 2, 3, padding1) self.out_conv nn.Conv2d(base_dim * 2, 3, 3, padding1) self.t_emb TimestepEmbedding(base_dim) self.gamma nn.Linear(base_dim, base_dim) def forward(self, x_t, y, t): x torch.cat([x_t, y], dim1) # 条件拼接 t_emb self.t_emb(t) h self.in_conv(x) w self.gamma(t_emb).view(-1, h.shape[1], 1, 1) h h * (1 w) # 时间步信息注入 h torch.relu(self.mid_conv(h)) return self.out_conv(h)这段代码有几个关键点。第一in_channels6是因为彩色图的x_t和退化图y各3个通道拼接后是6灰度恢复任务则改成2。第二网络输出通道仍是3因为预测目标是与输入图同尺寸的噪声图。第三这里用了一个简化U-Net骨架真正实验时要替换为带4次下采样/上采样和skip connection的结构时间步注入也可以改成scaleshift的AdaGN形式。判断模型是否写对可以用model(torch.randn(2,3,64,64), torch.randn(2,3,64,64), torch.randint(0,1000,(2,)))跑一次forward输出形状应为(2,3,64,64)。3.3 训练主循环loss计算、EMA、学习率、日志train.py主循环是完整可运行代码的核心。注意两个细节加噪的干净图来自clean目标退化图只作为条件loss是预测噪声和真实噪声的MSE。for step, (clean, degraded) in enumerate(dataloader): clean clean.to(device) degraded degraded.to(device) t torch.randint(0, num_timesteps, (clean.size(0),), devicedevice) x_t, noise forward_diffusion_sample(clean, t, alphas_cumprod) pred_noise model(x_t, degraded, t) loss torch.nn.functional.mse_loss(pred_noise, noise) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema_update(model, ema_model, decay0.999) if step % 50 0: print(fstep {step:6d} | loss {loss.item():.6f})训练时用x_t从clean加噪而不是从degraded加噪原因是目标噪声必须和输入严格对应模型学“把加噪后的clean还原”的同时参考degraded做条件约束。EMA维护一组更平滑的权重推理时用ema_model通常比model高0.2~0.5dB值得一开始就加上。clip_grad_norm_防止早期loss爆炸学习率一般用1e-4AdamW配合cosine schedule。把dataset.py、model.py、train.py、sample.py按上面的关键函数填入就是一份能跑通128×128高斯去噪的最小完整代码。4. 详细实验操作流程从数据准备到完整采样实验4.1 数据准备与退化模拟去噪为例实验流程第一步是造“成对数据”。用公开的BSD400或DIV2K数据集预先切成128×128的patch退化图现场加高斯噪声sigma25/255是常用配置。加噪与裁剪的示例代码import cv2 import numpy as np def add_gaussian_noise(img, sigma25): # img: float32值域 [0,1]sigma 以像素值为单位 noise np.random.randn(*img.shape).astype(np.float32) * (sigma / 255.0) return np.clip(img.astype(np.float32) noise, 0.0, 1.0) def random_crop_pair(clean, degraded, crop_size128): H, W clean.shape[:2] top np.random.randint(0, H - crop_size 1) left np.random.randint(0, W - crop_size 1) return (clean[top:topcrop_size, left:leftcrop_size], degraded[top:topcrop_size, left:leftcrop_size])sigma这里有个最容易踩的暗坑如果输入值域是[-1,1]加噪标准差应为sigma/127.5这是把噪声幅度相对值域做映射。很多第一次复现的人在这里翻车值域用[0,1]sigma却直接写25最后采样结果整体发灰。所以我建议dataset.py里统一走[-1,1]读图后执行img img / 127.5 - 1这和DDPM加噪公式天然一致。完整的数据集类按PyTorch惯例写读入图片、在线加噪、返回clean和degraded即可。退化模拟必须和你的恢复目标一致做去噪就加高斯噪声做超分就下采样再上采样做去模糊就要模糊核。训练时用什么退化参数测试时就用什么参数这是实验可复现的基础。4.2 训练配置与超参数一张表看懂最常用的起步值起步阶段的参数我一般直接照下表填参数推荐值说明num_timesteps1000DDPM标准设置显存紧张可降到500beta_schedulelinear 1e-4→0.02最常用噪声表高分辨率可换cosinecrop_size128兼顾纹理表达与显存占用batch_size812G显存可跑16G可提到16learning_rate1e-4AdamWcosine衰减ema_decay0.999推理权重用EMA版本grad_clip1.0防前期loss爆炸epochs100数据量小时建议数step而不是epoch这些值不是拍脑袋定的num_timesteps降到500时alphas_cumprod要重新生成采样步数也相应降到100左右crop_size不建议小于64否则纹理信息不够去噪结果容易糊batch_size和显存关系密切不够就减crop或开AMP。4.3 推理采样流程从纯噪声到恢复图的确定性采样训练完成后sample.py按“从噪声出发迭代去噪”的方式做恢复。下面这段是DDIM式更新用num_steps控制实际采样步数def sample_restore(model, degraded, num_steps200, eta0.0): model.eval() x torch.randn_like(degraded) # 从纯噪声图开始 with torch.no_grad(): for i in reversed(range(num_steps)): t i * (num_timesteps // num_steps) # 映射到训练步 t_batch torch.full((degraded.size(0),), t, devicedegraded.device) pred_noise model(x, degraded, t_batch) alpha_bar alphas_cumprod[t] alpha_bar_prev alphas_cumprod[t-1] if t 0 else 1.0 x0_pred (x - torch.sqrt(1 - alpha_bar) * pred_noise) / torch.sqrt(alpha_bar) x0_pred x0_pred.clamp(-1, 1) sigma_t eta * torch.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * torch.sqrt(1 - alpha_bar / alpha_bar_prev) x torch.sqrt(alpha_bar_prev) * x0_pred \ torch.sqrt(1 - alpha_bar_prev - sigma_t**2) * pred_noise \ sigma_t * torch.randn_like(x) return x0_pred这段代码用eta0的确定性更新每一步先用pred_noise反推出“干净图估计”x0_pred并clamp再按前向公式把这个估计重新加噪到下一步循环结束后返回最后的x0_pred。num_steps200是速度和质量比较稳的平衡点常见实验用100~200步想先验证pipeline可以跑100步。eta这个参数很关键eta0时没有随机噪声采样是确定性的多次推理结果一致方便调试eta0引入随机性可能带来更丰富的纹理但也会让结果不稳定。推理输入degraded必须做和训练完全一致的值域转换否则x0_pred第一步就会被clamp拉到错误区间。提示采样输出记得从[-1,1]映射回[0,1]再保存((x0_pred 1) / 2).clamp(0, 1)。这一步漏了后面算PSNR/SSIM全都会错位。5. 避坑与常见问题排查训练不收敛、图像发灰、采样发糊的五个现场5.1 loss在降采样图却“脏”先查条件拼接有没有配对错现象训练loss稳步下降采样结果却像“噪点色块”的混合体完全不是清晰的恢复图。原因最常见的是训练循环里把x_t误取自degraded而不是clean。目标噪声是“对clean加噪得到的噪声”输入x_t却是“对degraded加噪”两者语义不一致网络学的不是去噪而是噪声混淆。这种情况你用代码诊断插件去抓往往第一个暴露的就是数据流不对。解决把train循环里的forward_diffusion_sample(clean, t, alphas_cumprod)打出来核对再把model(x_t, degraded, t)改成model(x_t, torch.zeros_like(degraded), t)跑一次debug如果去掉条件还能学到基本去噪说明网络没问题问题一定在条件拼接。5.2 图像发灰、对比度低检查值域与clamp现象恢复图看起来“雾蒙蒙”细节不够锐利肉眼明显发灰但PSNR不算太低。原因大概率是图像值域和网络假设不一致。模型内部按[-1,1]设计代码里却把[0,1]的图直接喂进去或者采样时x0_pred被clamp到[-1,1]但前处理没有把图像归一到这个区间导致每一步的估计都系统性偏灰。解决固定统一值域。读图后立即img img / 127.5 - 1采样结束后out (out 1) / 2加噪噪声标准差按sigma/127.5计算。这个坑在换数据集时最容易复发建议直接写进dataset.py的模块注释里。5.3 训练到一半显存溢出或速度极慢现象batch16、crop128跑几十个step就OOM或者一晚上只跑完几千step明显不对劲。原因diffusion模型每步要同时维护x_t、noise、y和梯度显存占用比普通CNN高不少另外可能是没开混合精度cudnn benchmark也没开卷积反复autotune浪费大量时间。解决优先降crop到96或64batch减半然后开AMP用torch.autocast(device_typecuda, dtypetorch.float16)包forward/backward配合torch.amp.GradScaler。显存实在小就把una的base_dim从64降到32。注意力模块如果实现不当也会吃显存但用第3章那个简化骨架不会有这个问题。5.4 采样步数200步太慢减少后效果明显劣化现象200步出一张128×128图要一分多钟降成50步后图变糊或出现噪点。原因两个可能。一是t映射没对齐直接把reversed(range(num_steps))当成t输入网络而不是按比例映射到1000步空间二是更新公式里x0_pred没有正确clamp步数减少后误差累积。解决先确认映射是t i * (num_timesteps // num_steps)。然后对比100、150、200步的PSNR找到质量拐点。若必须快速采样再考虑DDIM、DPM-Solver或consistency蒸馏方向——那是后话别指望靠改一两个参数既快又好。5.5 训练集指标高、测试集崩盘数据量不够时的过拟合现象训练集上PSNR已经32dB以上测试集只有24~25dB两者差距巨大。原因扩散模型本质是生成模型数据量太小时会直接背训练集分布小数据集配上大batch这个问题会暴露得更快。解决用128×128随机裁剪把每张图变成几十个patch训练集用BSD400DIV2K混合并行加入随机翻转和旋转。如果条件允许先在无监督大规模噪声集合上预训练再在下游数据上微调。注意千万不要把测试集的裁剪图混进训练这是最常见的评测污染。6. 用固定测试集、固定随机种子做实验三个把结果做扎实的验证习惯先看固定随机种子的代码这是所有对比实验的地基import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False第一个习惯每次训练结束必须用同一个测试子集、同一个seed重新采样。这样不同配置之间的对比才是公平的。如果每次推理都放任随机源变化哪怕同一份权重两次采样结果也会有肉眼可见的差异指标对比毫无意义。第二个习惯评估时统一用skimage.metrics.peak_signal_noise_ratio和structural_similarity但一定要小心data_range参数。图像值域[-1,1]时data_range2.0值域[0,1]时data_range1.0。我第一次复现去噪实验时PSNR比baseline低0.5dB查了一天发现是data_range写成了1.0而模型输出是[-1,1]。这类“低分”往往不是模型问题是评测口径问题。第三个习惯每次实验记录轻量日志。我在每次训练前会把config.yaml内容、数据集版本、退化参数、最终PSNR/SSIM这四样复制进一个CSV同一配置至少跑两次保留seed记录。扩散模型采样有随机性一次性结果不可信。每次多花10分钟记录能少走一周弯路。另外可以固定一个“小验证集”训练每500步采样3张图肉眼观察恢复质量。loss曲线只能告诉你是否收敛不能告诉你恢复效果是否真实在屏幕上花20秒看的直观反馈比任何指标都早。我自己后来养成的习惯是任何新配置先不跑完整训练先在测试集上拿20步采样跑通一次全链路再放开训练等链路通了再加速调参。希望这个习惯也能帮到你。本文还有配套的精品资源点击获取