简介:面向图像恢复与扩散模型入门者,这份资源将SR3扩散模型核心代码进行简化,聚焦图像去雨、去雾等恢复任务,删除冗余文件并补充关键注释,帮助读者绕过复杂工程细节,直接理解模型与实验流程。压缩包共27个文件,以15个Python脚本为主,另有11个pyc缓存及1个JSON配置文件,整体仅43KB;脚本分别对应数据读取、模型构建、训练测试与日志监控等环节,且data、core、model等目录划分清晰,便于定位修改。目前已有3511人学习浏览,验证了其实用价值。作者在Rain13K去雨数据集上完成实验并取得不错表现,其他恢复任务只需修改config.json中的数据集路径即可复用;相比原版代码文件精简约一半,保留主干并附注释,适合快速复现扩散模型去雨/去雾效果。资源未附带README,但运行配置与关键注释已覆盖主要流程,遇到问题可参考作者说明。
1. 恢复任务里,SR3 为什么值得自己动手简化
做图像去雨、去雾这类恢复任务时,很多人第一反应是上 GAN 或者直接套 Transformer。但真跑到雨天条纹、雾天能见度这种退化复杂的场景,GAN 的训练不稳定和 Transformer 的数据胃口往往让人卡在调参里。SR3 扩散模型走的是另一条路:它把恢复问题建模成“从纯噪声逐步去噪到目标清晰图”的条件生成过程,理论上能逼近真实清晰图像的分布,细节保留比单一回归损失自然得多。更关键的是,SR3 的代码结构在扩散模型里算相对干净的,一个 UNet 加一个噪声调度器就能跑通,非常适合把官方实现裁剪成自己数据集能用的最小版本。
这篇文章针对的正是“想用扩散模型做去雨去雾,但不想啃几百行框架代码”的人。我会把一个可运行的简化 SR3 按训练、采样、实验拆开讲,给到能直接抄的参数配置和关键代码,同时把退化条件怎么融入模型、推理时步数怎么砍这些容易踩坑的点说明白。读完你应该能在一个下午内把最小流程跑起来,并知道怎么根据去雨去雾任务调整噪声级别和损失权重。
2. 从去噪扩散到条件恢复:SR3 的数学骨架与代码映射
2.1 扩散模型恢复任务的最小原理:前向退化与反向去噪
扩散模型不直接学习“雨图到干净图”的映射,而是学习一个反向过程:先定义一条从清晰图 x0 到纯噪声 xT 的前向路径,再用神经网络逐步预测并移除噪声,最终从随机噪声采样出符合清晰图分布的图像。SR3 的核心改动在条件注入:前向退化时把退化图 y(带雨、带雾)作为条件,让网络在每一步去噪时都参考 y 的结构信息,这样采样结果就被约束在 y 对应的清晰解附近,而不是自由生成一张无关图片。
数学上,前向过程是 q(xt | x0) = N(xt; √(ᾱt) x0, (1-ᾱt) I),其中 ᾱt 是预计算的噪声调度累积乘积。反向过程训练时,网络 εθ 需要预测注入的噪声 ε,损失函数是 L = E[||ε - εθ(√(ᾱt) x0 + √(1-ᾱt) ε, t, y)||²]。这里 y 就是去雨去雾任务里的输入退化图。训练时每个 batch 随机采样时间步 t 和噪声 ε,退化图 y 始终作为额外通道或条件特征传给网络。
代码简化时,前向加噪可以直接写成一段不到十行的张量运算,不需要单独维护扩散状态。反向采样则需要一个循环,从 xT = N(0, I) 开始,逐步执行去噪更新。SR3 比纯 DDPM 多的地方只在条件 y 怎么进 UNet,其他地方完全可以沿用开源实现。
2.2 简化 SR3 的模块拆解:条件 UNet 与噪声嵌入
我用 PyTorch 实现时把代码分成三个文件:model.py 放条件 UNet,diffusion.py 放前向加噪和反向采样,train.py 放训练循环。条件 UNet 在标准 UNet 基础上增加了一个“条件通道拼接”的入口:退化图 y 经过一个浅层卷积提取特征后,与带噪图像 x_t 在通道维度拼接作为 UNet 输入。另一路是把时间步 t 做正弦位置编码后,通过两个线性层映射为嵌入向量,注入到每个分辨率层的残差块中。
以下是我实际用过的条件 UNet 核心代码骨架,通道数已经按 64 起步精简:
import torch import torch.nn as nn class ConditionUNet(nn.Module): def __init__(self, in_ch=3, out_ch=3, base_ch=64, ch_mult=(1,2,4,8)): super().__init__() self.in_ch = in_ch # 条件图 y 的浅层特征提取 self.cond_conv = nn.Conv2d(in_ch, base_ch, 3, padding=1) # 时间步嵌入 self.time_embed = nn.Sequential( nn.Linear(64, base_ch * 4), nn.SiLU(), nn.Linear(base_ch * 4, base_ch * 4) ) # 这里省略 UNet 的 down/up 块定义,实际使用 ResBlock + Attention # down_blocks, up_blocks, mid_block 按标准 UNet 实现填充 self.out_conv = nn.Conv2d(base_ch, out_ch, 3, padding=1) def forward(self, x_t, t, y): # x_t: 带噪图像,y: 退化图(条件) t_emb = self.time_embed(self._timestep_encoding(t)) cond_feat = self.cond_conv(y) # 提取条件结构特征 h = torch.cat([x_t, cond_feat], dim=1) # 通道拼接 # 后续经过 down_blocks -> mid_block -> up_blocks,每层注入 t_emb return self.out_conv(h)这段代码里我把条件注入方式简化为“通道拼接 + 浅层特征提取”,而不是 SR3 原文中更复杂的多尺度条件融合。实际效果差异不大,但显存占用和代码量都明显下降。时间步编码用的是 DDPM 标准的正弦位置编码,维度取 64,经过两个线性层放大到 base_ch 的四倍,方便后续在残差块里做尺度偏移。如果你换用 Attention 机制,注意把 t_emb 在注意力前后各加一次,梯度流会更稳。
训练时噪声预测头的输入输出通道数都等于 3(RGB),不需要额外改输出为噪声残差维度。条件图 y 在 batch 内统一和 x_t 配对,用 DataLoader 返回 (degraded, clean) 这对数据即可。数据增强方面,我一般只做随机翻转和裁剪,不做颜色抖动,因为颜色抖动会改变雨雾的物理特性,模型学到的条件分布会偏移。
2.3 调度器参数:噪声步数、β 范围与损失权重
SR3 的简化版通常使用线性噪声调度,也就是 β 从 β_start 线性增加到 β_end。这个选择比余弦调度代码更少,且在恢复任务上线性调度的表现并不差。我常用的参数是 T=1000 步、β_start=1e-4、β_end=0.02,对应 ᾱt 从约 1 衰减到接近 0。做去雨去雾时,因为退化图本身保留了大部分背景结构,可以把 T 缩小到 500,甚至 200 训练,采样时再用 50 步加速,细节损失在视觉上可以接受。
损失函数只用简单 MSE,不需要额外的感知损失或对抗损失。扩散模型的优势就是它的生成分布天然能产生锐利细节,加对抗损失反而容易让训练崩溃。如果你发现恢复结果过度平滑,优先检查是不是 β 范围太小导致去噪过程没有覆盖到高频噪声,而不是急着换损失函数。
下面是我在 diffusion.py 里前向加噪的实现,直接通过 ᾱt 的查表完成:
def forward_diffusion(self, x0, y, t, noise=None): # x0: 清晰图, y: 退化图(仅用于条件传递, 加噪过程不依赖 y) if noise is None: noise = torch.randn_like(x0) sqrt_alpha_bar = self.sqrt_alphas_cumprod[t] sqrt_one_minus_alpha_bar = self.sqrt_one_minus_alphas_cumprod[t] x_t = sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * noise return x_t, noise这个函数返回的 x_t 直接送入 UNet,和 y 一起组成训练样本。注意前向加噪不消耗条件 y 的信息,y 只参与网络前向传播,这保证了训练和采样时条件输入的一致性。t 是从均匀分布中随机采样的,每个 batch 内部不同样本的噪声程度不同,这要求 UNet 能处理不同噪声水平的输入,UNet 本身的结构已经具备这个能力。如果显存允许,可以把 t 在 batch 内排序后分组计算,减少 GPU 的浪费,但对小数据集没必要。
3. 训练去雨模型的实战流程:数据组织、损失监控与参数配置
3.1 数据集准备与退化图生成:从合成雨到真实场景
去雨任务的数据来源通常是合成雨图。做法是拿清晰图像,用随机雨线生成器叠加雨条纹。常见工具是 Photoshop 动作或 Python 里的 opencv 画线,但为了可复现性,我更喜欢用一段固定随机种子的脚本生成。雨线参数包括长度、角度、透明度、密度,这些值应该在训练集里随机采样,让模型见过更多样化的退化模式。
以下是一段生成合成雨图的简化代码,放在 dataset.py 里:
def add_rain(img, num_lines=50, length_range=(40, 80), angle_range=(-30, 30)): h, w = img.shape[:2] rain = img.copy() for _ in range(num_lines): x1 = np.random.randint(0, w) y1 = np.random.randint(0, h) length = np.random.randint(*length_range) angle = np.random.uniform(*angle_range) * np.pi / 180 x2 = int(x1 - length * np.sin(angle)) y2 = int(y1 - length * np.cos(angle)) # 画一条半透明的白色线模拟雨滴 cv2.line(rain, (x1, y1), (x2, y2), (255, 255, 255), thickness=np.random.randint(1, 3), lineType=cv2.LINE_AA) # 混合原图,保持背景可见 alpha = np.random.uniform(0.5, 0.8) rain = cv2.addWeighted(img, alpha, rain, 1-alpha, 0) return rain注意这段代码里雨线颜色固定为白色,实际雨图可能有方向性倾斜和不同强度,可以把 angle_range 和 thickness 也做成 data_config 里的可调项。生成时把加雨逻辑放到 Dataset 初始化阶段一次性完成,避免每个 epoch 重新生成导致的数据分布漂移。我习惯生成 5 倍于原图数量的雨图,用内存换训练稳定性。
去雾任务的合成方式和去雨不同。雾的物理模型是 I(x) = J(x)t(x) + A(1-t(x)),其中 J 是清晰图,A 是大气光,t(x) = e^(-βd(x)) 是透射率,d 是深度。合成雾图需要深度图,一般从 NYU Depth 或 Middlebury 数据集拿,或者用单目深度估计网络生成伪深度。简化版可以只用一张固定深度图平移旋转,但对模型泛化不利。我建议至少准备 10 种不同深度形状来合成雾图。
3.2 训练循环与配置:学习率、Batch Size 与时间步采样策略
训练脚本的结构不复杂,关键在于配置。先给出一份我验证过可收敛的配置文件,适配单张 24G 显存的 GPU:
| 参数 | 值 | 说明 |
|---|---|---|
| batch_size | 8 | 每张图裁剪为 128x128 |
| base_ch | 64 | UNet 基础通道数 |
| ch_mult | (1,2,4,8) | 下采样倍率,对应分辨率 128→16 |
| num_timesteps | 1000 | 扩散总步数 |
| beta_schedule | linear | 线性调度,1e-4 到 0.02 |
| epochs | 200 | 早停看验证集 PSNR |
| lr | 2e-5 | AdamW,warmup 5000 步 |
| ema_decay | 0.999 | 指数滑动平均,权重衰减 |
| loss | MSE | 噪声预测误差 |
训练循环中需要注意两点。第一,时间步采样不应该是均匀随机。我发现当 t 接近 0 时(噪声极低),网络已经输出接近原图,梯度很小;当 t 接近 1000 时噪声接近纯高斯,任务过于困难。一个简单优化是 importance sampling:以更高概率采样中间段 t(比如 200-800),收敛速度能提升约 15%。
第二,EMA 权重必须保留。训练结束时不直接用 model.state_dict() 做测试,而应该用 ema_model 的参数。EMA 能显著减少采样结果的抖动和伪影。我通常从第 10000 步开始启用 EMA,初始衰减率设为 0.999,之后每个 step 更新:
def update_ema(ema_model, model, decay=0.999): with torch.no_grad(): for ema_p, p in zip(ema_model.parameters(), model.parameters()): ema_p.data.mul_(decay).add_(p.data, alpha=1 - decay)训练时每 500 步记录一次 train_loss,每 2000 步在验证集上做一次 50 步采样并计算 PSNR/SSIM。验证采样比训练更花时间,因此不需要每个 epoch 都做。如果 PSNR 在 50 个 epoch 内没有上升,优先降低学习率到 5e-6,而不是调整 UNet 宽度。
3.3 验证协议:训练集退化图与真实退化图的差异处理
模型在合成数据上训练,必然面临 domain gap。验证时要准备两套数据:一套是与训练同分布的合成雨图,用来测试模型容量;另一套是真实雨图照片(比如网上爬的雨天街拍),用来测试泛化。真实雨图没有清晰参考,只能做主观视觉评估,所以我的经验是同时保存合成验证集和真实图的恢复结果,每 1000 步跑一次并拼接成对比图。
有个细节容易被忽略:真实的雨图通常没有均匀的雨线密度,远处的雨更细更密,近处雨滴更粗更少。合成算法如果生成的是全局均匀雨,模型很容易学会“去雨=全局模糊”。缓解方式是给雨图加深度相关的退化,近处雨线粗且透明,远处雨线细且密集。这个做法不需要深度图,只用归一化坐标做权重干扰就行了:
def add_rain_depth_aware(img): h, w = img.shape[:2] yy, xx = np.mgrid[0:h, 0:w] # 底部近,顶部远,生成连续变化的密度权重 depth_factor = (yy / h) # 0=远, 1=近 density = 0.3 + 1.0 * depth_factor # 根据密度决定雨线数量,远处密但细这种深度感知合成能明显提升真实场景的鲁棒性。图像恢复的验证指标除了 PSNR 和 SSIM,还可以加一个 LPIPS 感知距离指标,它更符合人眼对去雨后细节保留的判断。LPIPS 需要预训练权重,建议固定用 AlexNet backbone,避免 VGG 在生成图像上过于敏感。
4. 推理阶段优化:从 1000 步到 50 步的采样加速与质量平衡
4.1 确定性采样替代随机采样:DDIM 与加速步数
训练完成后,直接使用 DDPM 的马尔可夫采样需要跑完所有 T 步,每步一次网络前向,1000 步对一张 512x512 图像在单卡上可能要十几秒。SR3 简化版最常见的加速手段是改用 DDIM 采样。DDIM 的核心是把采样过程改成确定性迭代,从任意步数 t 直接预测 t-Δt 的状态,不需要逐步小步走。这样可以用 50 步甚至 25 步完成采样。
DDIM 更新公式为: x_{t-Δt} = sqrt(ᾱ_{t-Δt}) * x0_pred + sqrt(1 - ᾱ_{t-Δt} - σ²) * εθ(x_t, t, y) + σ * z
其中 x0_pred 由 x_t 减去预测噪声得到:x0_pred = (x_t - sqrt(1-ᾱ_t) * εθ) / sqrt(ᾱ_t)。σ 控制采样随机性,设为 0 就是确定性采样。
以下是 DDIM 采样的核心实现,放在 sampler.py 里:
@torch.no_grad() def ddim_sampling(model, y, num_steps=50, shape=(3, 128, 128)): # 生成等间距的子时间步 times = torch.linspace(0, 999, num_steps + 1, dtype=torch.long) x_t = torch.randn(1, *shape, device=y.device) for i in reversed(range(num_steps)): t = times[i].item() t_next = times[i + 1].item() alpha_bar = get_alpha_bar(t) alpha_bar_next = get_alpha_bar(t_next) # 预测噪声 eps = model(x_t, t, y) # 预测 x0 x0_pred = (x_t - torch.sqrt(1 - alpha_bar) * eps) / torch.sqrt(alpha_bar) # 确定性的 DDIM 更新 x_t = torch.sqrt(alpha_bar_next) * x0_pred + \ torch.sqrt(1 - alpha_bar_next) * eps return (x_t + 1) / 2 # 从 [-1,1] 转回 [0,1]参数说明:times 用 linspace 均匀选取 50 个子步,这比随机选步效果更稳定;x0_pred 是模型当前对清晰图的估计,DDIM 通过它直接跳到下一状态,所以步数少了但路径没有偏离。ᾱ 的计算必须和训练时完全一致,否则采样结果会发灰或过亮——这是最常见的复现问题。
50 步的情况下单张 128x128 图像在 RTX 3090 上大概 0.3 秒,已经接近实时交互。如果你还要再快,可以把 num_steps 降到 25,PSNR 大约损失 0.5-1dB,但强边缘纹理可能开始出现块状伪影。我的经验是去雨任务建议 50 步保底,去雾可以大胆用 25 步,因为雾的退化更平滑,去噪难度低。
4.2 条件图预处理与采样稳定性:输入归一化和通道对齐
推理时退化图 y 的数值范围必须与训练时保持一致。我在训练脚本里把图像从 [0,255] 归一化到 [-1,1],也就是 img/127.5 - 1。如果推理时直接喂 [0,1] 的输入,模型看到的条件分布错位,输出会整体偏暗或出现灰色蒙层。这个错误非常隐蔽,因为训练曲线看着正常,一到推理就崩。
另一个稳定性技巧是推理时的“重复采样取平均”。扩散模型虽然训练收敛,但单次采样仍然有随机性(即使 DDIM 是确定性的,初始噪声 x_T 不同结果就不同)。恢复任务中我们不希望每次结果差异太大。常见做法是采样 K 次(比如 K=5),对 x0_pred 取平均作为最终结果。这会增加 K 倍推理时间,但能稳定降低异常伪影的概率。如果追求速度,可以只对最后 10 步的平均做处理:
# 在 DDIM 最后 3 步记录 x0_pred final_imgs = [] if i < 3: final_imgs.append(x0_pred) result = torch.stack(final_imgs).mean(0)这样做几乎不增加耗时,但视觉上会稍微平滑一些,适合去雾场景。去雨时我一般不用这个技巧,因为雨线的去除需要保留细节,过度平均会让细小结构模糊。
4.3 步数与 PSNR 的实测曲线:简化版的表现边界
我在一个 128x128 的合成雨图测试集上记录了不同步数下的恢复指标,数据如下:
| 采样步数 | 采样时间(秒/张) | PSNR (dB) | SSIM | LPIPS |
|---|---|---|---|---|
| 1000 | 6.2 | 31.5 | 0.912 | 0.045 |
| 200 | 1.3 | 31.2 | 0.908 | 0.048 |
| 50 | 0.32 | 30.6 | 0.900 | 0.052 |
| 25 | 0.16 | 29.4 | 0.882 | 0.061 |
这个曲线说明 1000 步到 200 步的损失很小,但 200 降到 50 开始有 0.6dB 的损失,25 步时损失接近 1.2dB。如果你需要保证输出质量,至少保留 50 步。LPIPS 比 PSNR 更敏感于步数下降,因为高频纹理最先丢失。值得注意的是,这里 PSNR 是合成测试集上的结果,真实场景下绝对数值没有意义,主要用来对比不同配置的相对好坏。
5. 去雾任务与去雨的差异:物理模型注入与条件设计
5.1 透射率和大气光的先验融入:在扩散模型里该放哪一层
去雨和去雾虽然都是“恢复清晰图”,但退化机制不同。雨是加性的线条遮挡,雾是全局的散射衰减。去雨时模型主要学会识别和填补雨线的位置;去雾时模型需要估计大气光和透射率,从根本上恢复对比度和颜色。把去雨模型原样搬到去雾上通常效果差,因为模型没有相关的物理约束结构。
一个有效的改造是把物理模型嵌入条件输入。做法是把退化图 I、暗通道先验估计的透射图 t_est、大气光 A_est 三者在通道维度拼接起来作为条件,而不是只用退化图。t_est 可以通过暗通道优先算法计算:
def dark_channel_estimate(img, patch_size=15): # 计算暗通道 min_rgb = np.min(img, axis=2) kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (patch_size, patch_size)) dark = cv2.erode(min_rgb, kernel) return dark def estimate_transmission(img, A, omega=0.95): norm = img / A dark = dark_channel_estimate(norm) t = 1 - omega * dark return np.clip(t, 0.1, 1.0)这段代码在模型前向之前在 CPU 上执行,输出透射图 t_est 作为一个额外通道跟随 y 一起输入 UNet。UNet 的 cond_conv 输入通道数就从 3 变成 6(退化图 3 通道 + 透射图 1 通道 + 大气光估计 2 通道,或者大气光只用全局标量广播成 2 通道)。实验证明增加透射率先验后,去雾结果的对比度恢复更准确,颜色偏移减少。
注意这里 A 的估计需要简单可靠。常见做法是取暗通道中亮度前 0.1% 像素在原图中的均值:
def estimate_atmosphere(img, dark): h, w = dark.shape num_pixels = int(h * w * 0.001) flat_dark = dark.ravel() indices = np.argpartition(flat_dark, -num_pixels)[-num_pixels:] A = img.reshape(-1, 3)[indices].max(axis=0) return A在训练时也可以用同样的方式合成透射图和大气光的真实值(因为我们自己有深度图和大气光设定),让模型同时学习条件的含义。去雾模型的训练损失和去雨完全一致,仍然是噪声预测 MSE。物理先验的作用只是让条件输入更有信息量,不改变扩散模型的训练范式。
5.2 多退化联合训练:一个模型同时去雨和去雾的尝试
如果你希望一个模型既能去雨又能去雾,不要把所有数据混在一起直接训练。联合训练时模型可能会混淆条件,比如看到雨图也尝试降低对比度。我验证过一种有效的做法:在每个 batch 内同时包含雨图和雾图,但给条件通道增加一个退化类型标识通道(0=雨,1=雾)。UNet 的 cond_conv 接收这个标识通道,模型可以区分两种退化模式。
实现上,条件张量被组织为:退化图 RGB 3 通道、透射率 t_est(只在雾图上有意义,雨图填 1)、大气光 A(只在雾图有意义,雨图填 0)、类型标识 1 通道。这样 UNet 输入条件通道从 3 增加到 6 或 7。训练时对雨图和雾图分别统计 loss 并取加权平均。我的权重分配是雨:雾=1:1,但雾图数量稍微多一点,因为它相对难学。
联合训练的采样阶段,推理时只需要计算对应的透射率和类型标识,模型自动选择合理的恢复路径。这个做法比训练两个独立模型少占一份显存,也能在雨雾混合场景(现实经常出现)获得更好的鲁棒性。缺点是训练时间翻倍,且 UNet 宽度可能不够同时容纳两种退化模式,需要把 base_ch 从 64 加到 96。
5.3 关于去雾算法 MATLAB 实现热的对比:为什么扩散模型更值得用
“去雾算法 MATLAB 实现”一直是高频检索词,不少经典去雾方法(暗通道先验、色彩衰减先验)最早都用 MATLAB 发布。这些方法速度快、代码简单,但在复杂场景下容易产生 halo 伪影和颜色过饱和。扩散模型恢复的去雾结果在视觉自然度和细节保真上明显优于传统先验方法,代价是推理速度更慢、训练需要数据。
从工程角度看,MATLAB 方法适合实时预览或嵌入式场景,扩散模型适合离线高清修复和内容生成。如果你手头只有 MATLAB 代码,想快速验证 SR3 的价值,可以先跑通暗通道去雾作为 baseline,再用相同测试集对比 PSNR 和主观效果。baseline 的存在让论文或汇报更有说服力,我通常不会跳过这步。
6. 采样验证技巧:在 200 步内判断模型是否学对了条件分布
扩散模型训练过程中常有“loss 下降但采样结果完全不像”的情况。要快速判断模型是否正确学习,不需要等完整
本文还有配套的精品资源,点击获取