简介面向图像恢复任务研究者与进阶开发者本资源基于SR3扩散模型进行代码精简与重构重点解决去雨、去雾等恢复任务的落地调试难题。源码对原始项目做了大幅瘦身删除冗余小文件与无用代码块补充关键注释并在Rain13K去雨数据集上完成实验验证效果良好其他恢复任务通常只需修改配置文件中的数据集路径即可直接使用。资源包共27个文件以15个Python脚本为核心涵盖数据加载、网络模型、训练测试、指标与日志模块另有少量编译缓存文件和1个JSON配置文件整体仅43KB便于快速阅读与迁移。已有3511人学习浏览适合具备一定深度学习和PyTorch基础、希望快速上手扩散模型做图像恢复实验的读者。由于包内不含README作者提供邮箱联系方式可私信获取使用说明或进一步交流。1. 为什么我决定对一个SR3扩散模型动手瘦身一个周末我把原版SR3扩散模型仓库拉下来准备处理图像去雨。真正跑起来才意识到麻烦不在扩散模型本身而在于散落在core/model/sr3_modules里的几十个小文件光验证入口就有两个logger还分版本更不用说一堆可选模块把训练路径绕得云里雾里。我砍掉近一半无用代码保留DDPM核心、UNet主干和数据加载在Rain13K这个被MPRNet、Restormer反复使用的合成雨图数据集上重跑得到的去雨结果和迭代稳定度都够用。这份简化版SR3_IR就是这次瘦身实战的记录面向想拿扩散模型直接改业务场景的人讲哪些文件要保留、哪些参数会左右训练结果以及验证时会踩的坑。2. SR3扩散模型原理与代码瘦身关键2.1 扩散模型在图像恢复中的位置SR3本质上是将DDPM那套“正向加噪、逆向去噪”框架迁移到图像恢复任务并非只能做超分。正向过程是给定一张干净图像 x_0在 T 步内逐步注入随机高斯噪声得到越来越接近纯噪声的 x_t逆向过程则训练一个UNet网络输入带噪图像和任务条件预测当前步所叠加的噪声从而逐步还原出干净图像。对于图像去雨来说雨水条纹可以看作叠加在背景上的一种结构化噪声扩散模型学到的实际是“给定有雨图条件时从带噪图还原干净图”的条件分布。这比直接用L2损失回归一个去雨结果更接近生成式建模能保留更多高频纹理。和MPRNet、Restormer这类确定性深度网络相比SR3的代价是迭代步数多、训练显存占用高收益是对复杂场景下的雨痕形态更鲁棒生成的背景细节也更自然。理解这点之后对原版代码做删除和精简时就不会误伤核心采样逻辑。2.2 原版代码里哪些是非必要的原仓库里有很多为通用实验设计的多余模块不同损失函数接口、多尺度特征融合分支、分布式训练辅助逻辑、两个验证入口。这些文件对初学者不友好把训练入口和模型定义拆得非常分散一旦报错栈信息跳来跳去很难定位到具体张量 shape 不匹配的位置。我按“留一个进程、留一个损失、留一个验证入口”的原则清理去掉的文件包括重复的 eval 脚本、可选的下采样模块、多种训练策略的 backbone 封装、非必要的可视化回调。保留的核心部分如下原文件/模块简化后处理原因多个 train 入口合并为main.py只保留单进程训练流程多个验证脚本合并进test.py减少入口才能集中调试多损失选择器固定使用 L1 Loss去雨任务 L1 足够稳定各种 distributed 工具删除单卡调试更简单logger 代码精简到core/logger.py只记录 loss、PSNR、模型权重2.3 简化后的SR3_IR目录与文件职责最终目录结构和原始仓相比少了一半文件每个文件的功能非常清楚SR3_IR/ ├── main.py ├── test.py ├── data/ │ ├── __init__.py │ ├── dataset.py │ └── util.py ├── core/ │ ├── metrics.py │ └── logger.py ├── model/ │ ├── __init__.py │ ├── base_model.py │ ├── model.py │ ├── networks.py │ └── sr3_modules/ │ ├── __init__.py │ └── ddpm_modules.py └── config/ └── config.jsonmain.py负责整个训练循环test.py负责验证和结果输出data/dataset.py里只留了Rain13K的读取与裁剪逻辑data/util.py放图像归一化、随机翻转这类工具函数model/networks.py是UNet结构model/sr3_modules/ddpm_modules.py放扩散前向加噪、采样和beta调度config/config.json集中控制数据集路径、模型参数、训练超参数。这样做的好处是当你改一个batch size或者想要换数据集路径时不再需要在一堆Python文件之间跳转所有可能影响运行结果的开关都在同一个配置文件里。3. Rain13K上去雨训练配置、数据与执行流程3.1 把SR3从超分改成去雨的数据接口原版SR3的数据集接口读入的是低分辨率图和高分辨率图做去雨时需要把输入改成“有雨图”和“干净图”。我在dataset.py中重新定义了一个配对数据类核心是让数据集目录下同时存在 rain 和 clean 两个子目录图片按相同文件名对应。Rain13K 通常提供 13000 对室内外合成雨图其中训练集划分为 12000 对验证集 1000 对。常见做法是直接用torchvision.transforms做随机裁剪和翻转但要注意扩散模型对尺寸非常敏感。我的版本里默认把配对图随机裁剪成 64x64 的 patch这个尺寸在 4 张 batch size 下单卡能勉强跑起来。如果不先裁到小尺寸原图 512x512 几轮迭代就会把显存撑爆。需要说明的是扩散模型训练时并不需要特别复杂的在线数据增强随机水平翻转和随机旋转90度已经足够。重点是把图像从 [0,1] 转换到 [-1,1]因为DDPM的正向加噪和逆向采样的目标空间都围绕这个范围设计。3.2 config.json中的关键参数我把原先分散在Python里的超参数全部收拢到config/config.json整个去雨实验只需要改这个文件{ data: { train: { rain_dir: data/Rain13K/train/rain, clean_dir: data/Rain13K/train/clean }, eval: { rain_dir: data/Rain13K/test/rain, clean_dir: data/Rain13K/test/clean } }, model: { in_channel: 3, out_channel: 3, inner_channel: 64, norm_groups: 32, channel_multiplier: [1, 2, 4, 4], with_attn: [true, false, false, false], timesteps: 1000, beta_start: 0.0001, beta_end: 0.02 }, train: { batch_size: 4, patch_size: 64, epochs: 100, lr: 2e-5, save_dir: checkpoints/ } }这里timesteps是扩散过程的总步数1000 是DDPM原论文的标准设定beta_start和beta_end控制噪声调度范围数值过大会导致训练初期噪声过强。inner_channel64决定了UNet第一层卷积通道数如果显存充足改成 128 能提升容量但训练时间也接近翻倍。norm_groups32是每组归一化的通道数需要注意 batch size 必须能被norm_groups整除否则在 GroupNorm 层会报尺寸错误。3.3 训练循环正向加噪与条件输入main.py中的核心训练循环可以看成四步取一对图像、随机采样时间步、正向加噪、让模型预测噪声。简化后的代码保留了原版最关键的结构同时补充了shape注释def forward_process(clean, t, alpha_bar): # clean: [B, 3, H, W], 取值范围[-1, 1] noise torch.randn_like(clean) # alpha_bar 是累计噪声调度参数从config的beta序列计算得到 noisy torch.sqrt(alpha_bar[t])[:, None, None, None] * clean \ torch.sqrt(1 - alpha_bar[t])[:, None, None, None] * noise return noisy, noise for epoch in range(epochs): for step, (rain, clean) in enumerate(train_loader): # rain: 有雨图条件, clean: 干净图 t torch.randint(0, config[model][timesteps], (rain.size(0),), devicedevice).long() noisy_clean, noise forward_process(clean, t, alpha_bar) pred_noise model(rain, noisy_clean, t) loss F.l1_loss(pred_noise, noise) optim.zero_grad() loss.backward() optim.step()注意model的第一个输入是rain第二个输入是noisy_clean这和原始SR3里“低分辨率图带噪高分辨率图”的条件方式完全等价。模型不是直接预测干净图而是预测当前步的噪声原因在于DDPM的训练目标在噪声空间里更平滑网络更容易收敛。训练时用的alpha_bar是根据beta_start和beta_end预先算好的累计参数随着t增大alpha_bar[t]越来越接近 0noisy_clean就越来越接近纯高斯噪声。3.4 监控、模型保存与启动命令core/metrics.py里保存了PSNR和SSIM的独立实现验证阶段调用test.py对验证集做完整采样。训练过程中core/logger.py会记录每个epoch的平均loss、每固定步数输出的PSNR趋势以及最后一轮生成的去雨demo图。整个实验启动非常直接python main.py -c config/config.json如果只是想验证数据读取和模型尺寸没问题可以把epochs临时改成 1并把batch_size减小到 1。这里-c参数在main.py中用argparse解析配置文件路径必须和当前工作目录对应或者写绝对路径。实际运行中损失日志会以字符串追加到log/train.log一旦崩溃优先看这个文件里最后几条信息。4. SR3去雨实现ddpm_modules、UNet与可读注释4.1 预定义扩散过程beta调度与采样model/sr3_modules/ddpm_modules.py是整个扩散模型的核心。它没有复杂网络结构只做一件事根据beta_start和beta_end生成线性beta序列并据此计算前向加噪所用的alpha_bar。这段代码虽然是原样保留但我为每个公式补了注释重点说明张量维度和取值范围betas torch.linspace(beta_start, beta_end, timesteps).float() # betas: [timesteps], 每个元素代表该步加入的噪声方差 alphas 1.0 - betas alpha_bar torch.cumprod(alphas, dim0) # alpha_bar[t] 是0到t步累计衰减系数取值从0.99逐渐降到接近0这里最容易被忽略的是alpha_bar需要放在 GPU 上并且在计算alpha_bar[t]时要展开成[B, 1, 1, 1]才能和图像张量进行广播乘法。如果忘记形状对齐会在训练中途报出非常奇怪的 broadcast error。还需要注意的是当timesteps很大时betas序列的末端已经非常接近 0但这不意味着要修改调度。线性调度虽然简单却能让训练稳定换成 cosine 调度会让采样步数减少但对去雨任务的提升并不明显。4.2 UNet条件网络的结构model/networks.py定义的UNet同时接收两个输入有雨图和带噪图。在我的实现里两个输入会在第一层卷积后合并到同一通道维度再进入下采样阶段class SR3UNet(nn.Module): def __init__(self, in_ch, out_ch, inner_ch, channel_multiplier, with_attn): # in_ch3表示带噪干净图通道condition_ch3表示有雨图通道 self.head nn.Conv2d(in_ch condition_ch, inner_ch, kernel_size3, padding1) # ... 后续是下采样、中间层、上采样参考DDPM官方结构 def forward(self, rain, noisy_clean, t): # rain: 有雨图条件noisy_clean: 带噪图 h self.head(torch.cat([rain, noisy_clean], dim1)) t_emb self.time_mlp(t) # 时间嵌入t_emb: [B, inner_ch], 会在每个尺度上通过add操作注入 return self.decode(self.encode(h, t_emb), t_emb)如果你要调整去雨效果优先改channel_multiplier和with_attn。例如channel_multiplier: [1,2,4,8]会明显增加UNet下最深层感受野有助于去除大范围雾霾但显存占用也会涨不少with_attn通常只在前两层保留自注意力因为全分辨率注意力计算量太大。4.3 注释写在哪里更有价值我不建议给所有文件都逐行加注释那样反而会让代码更冗长。简化版SR3里我主要给三类位置加了注释第一类是张量 shape 变化的节点例如torch.cat前后、reshape前后第二类是容易被改坏的超参数比如norm_groups和 batch size 的整除关系第三类是扩散公式中数值稳定性的说明例如torch.sqrt(alpha_bar[t])前的数值范围。IDE 的方法注释模板只能生成文档框架真正有用的还是手动写清楚每个张量的含义和维度。例如在 UNet 的下采样模块里我会写这样的注释# 输入 h: [B, C, H, W] # 经过下采样卷积后变为 [B, C*2, H/2, W/2] # 这里必须同步修改后续跳跃连接的channel数否则上采样拼接会失败这类注释的价值在于当你想把图像尺寸从64改到128时能一眼看出哪些位置需要跟着调整。很多初学者改patch size后训练崩溃就是因为没有注意到下采样次数和注意力层尺寸的耦合关系。5. PSNR验证、DDIM加速与常见坑5.1 去雨结果的验证方法与指标测试阶段不能像训练那样直接预测噪声需要从纯噪声开始逐步采样test.py中的验证脚本会自动加载训练好的权重并生成去雨图python test.py --ckpt checkpoints/best.pth \ --config config/config.json \ --output results/验证时默认对每张测试图做完整1000步采样这一步在CPU上会非常慢建议在GPU上执行。core/metrics.py中计算PSNR时会把范围从 [-1,1] 反归一化回 [0,255]避免因尺度不同导致PSNR虚高。每张图会生成三个文件有雨输入、模型输出、干净参考图方便逐张对比雨水纹理的消除程度。如果你觉得1000步太慢可以在测试阶段临时将timesteps设成 200并且用简单跳步采样。这种偷懒方式会让输出稍模糊但能用来快速筛选权重。5.2 训练不收敛与生成空洞的排查表实际操作里最容易遇到的是下面几类问题我整理了一张速查表现象可能原因检查方式loss 不下降学习率过高或 t 采样严重偏向高噪声区打印每步 alpha_bar 的范围把 lr 降到 1e-5采样结果全黑模型输出没有经过反归一化确认torch.clamp(img, -1, 1)后再转图像雨痕被去除但背景模糊训练 patch size 太小UNet 下采样后空间信息丢失将 patch_size 提高到 128观察验证PSNR大面积伪影条件图与带噪图在 cat 时通道顺序反了用固定种子跑一次 forward检查rain和noisy_clean的标准差显存不足norm_groups过大或 batch size 过大减小norm_groups到 8 或 16而非只降低 batch size其中最有迷惑性的是“训练loss缓慢下降但生成图一直有雾状残留”。这通常不是因为模型结构有问题而是beta_end设置过大导致早期噪声分布已经超出UNet能恢复的范围。把beta_end从 0.02 降到 0.01往往能明显改善最终输出。5.3 把同一套代码迁移到图像去雾因为数据接口已经在dataset.py中解耦迁移到图像去雾只需要替换数据集路径和修正图像预处理。雾图相比雨图往往是全局均匀的干扰条件特征更需要全局感受野所以要把channel_multiplier调深并把with_attn至少保留在第二个下采样层。我一般会直接建立一个新的数据集子目录让config.json里的rain_dir指向雾图目录clean_dir指向对应无雾图然后把模型输入输出通道从3改成3保持不变。最后验证时建议从测试集中挑出若干张纹理最复杂的样本单独跑一遍采样并且保存中间步骤的生成结果用来判断是在哪一步引入的雾状残影。如果从头到尾都带一层白雾就先检查beta调度和归一化范围如果只在最后几步出现就重新训练模型。这样做虽然略显笨拙但比反复调lr管用得多。本文还有配套的精品资源点击获取