ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

RIDNet图像去噪模型复现:从EAM模块搭建到PSNR/SSIM评估实战

2026/8/30 6:19:16 拓冰建站 浏览量
RIDNet图像去噪模型复现:从EAM模块搭建到PSNR/SSIM评估实战 简介本资源是基于PyTorch实现的RIDNet图像去噪模型完整复现项目面向深度学习初学者与计算机视觉方向实践者聚焦真实场景下的噪声抑制任务可直接用于单张图像去噪推理及PSNR/SSIM指标量化评估。压缩包共20个文件包含7个核心Python脚本如RIDNet.py模型定义、train.py训练入口、predict.py单图测试、test_benchmark.py批量评测、4个XML配置文件IDE与项目环境相关、3个pyc缓存文件及README说明文档等整体体积16.38MB结构清晰模块职责明确datasets管理数据、weights存放预训练模型、data提供示例图像、utils与loader封装通用工具与数据加载逻辑。已有539人学习下载用户可开箱即用——不仅获得已收敛的checkpoint.pth.tar模型还配套加噪测试、可视化预测、指标计算与结果保存全流程代码显著降低复现实验门槛。 图像去噪这个方向我前前后后折腾了大半年。RIDNetReal Image Denoising Network是ICCV 2019上的一篇真实图像去噪论文核心思路是用轻量级的EAM注意力模块配合全局残差学习在SIDD、DND这类真实噪声数据集上拿到过非常能打的效果。这篇博文就是我从零复现RIDNet的完整记录网络结构怎么搭、训练过程怎么配置、PSNR/SSIM指标怎么算、模型文件怎么加载全部摊开讲清楚。跑通之后直接拿真实图像测去噪效果肉眼可见不是那种训练完就吃灰的项目。如果你是刚接触图像去噪的学生或者论文写完了急需一份能跑的baseline又或者只是想把一张满是噪点的照片清理干净这份复现都能直接拿来用。我在GitHub上整理好了完整的工程结构包含训练脚本、评估脚本、预训练模型权重和推理demo下载下来按步骤执行就能出结果。1. 项目定位为什么复现RIDNet而不是其他去噪模型1.1 RIDNet在图像去噪里的位置图像去噪的深度学习路线基本是沿着两条主线走的。一条是CNN堆叠路线从DnCNN开始用残差学习预测噪声映射另一条是Transformer路线近两年比如Restormer、SwinIR这些靠全局注意力建模长距离依赖。RIDNet恰好站在中间它没有用复杂度爆炸的多尺度结构也没有引入自注意力的大矩阵计算而是通过一个设计极轻量的EAMEfficient Attention Module模块把通道注意力嵌入残差块在性能和计算量之间取了一个很好的平衡点。我选它复现的原因很简单第一模型结构干净核心组件少适合做二次开发和消融实验第二它在真实图像噪声场景下的泛化能力不错尤其是对SIDD这类传感器噪声数据的处理比纯合成噪声训练的DnCNN要稳第三相比动不动就几十层上百层的网络RIDNet的训练成本低得多单张消费级显卡就能跑完标准流程。1.2 项目包含哪些东西这个复现工程不是只有一份模型定义文件而是把从训练到评估再到推理的整条链路补齐了完整的RIDNet模型定义包括EAM模块、残差块、特征提取和重建分支。基于SIDD数据集的训练脚本包含数据加载、patch裁剪、数据增强、学习率调整策略。PSNR和SSIM的计算脚本支持灰度图和RGB图支持Y通道评估。训练好的模型权重文件.pth格式加载后可以直接对真实图像推理。一个简单的推理demo输入一张带噪图像输出去噪结果并打印PSNR/SSIM。整个工程用PyTorch实现不依赖任何第三方去噪库环境搭好之后import部分只涉及torch、torchvision、numpy、opencv这几个常规库。2. 环境搭建PyTorch版本选择与依赖安装2.1 PyTorch版本与CUDA匹配问题这次复现我用的环境是Python 3.9 PyTorch 1.13.1 CUDA 11.7。其实RIDNet的网络结构非常简单理论上PyTorch 1.x到2.x都能跑通但我建议不要用太新的版本原因后面在模型加载部分会讲到。如果是从零开始装环境推荐直接用conda创建独立环境conda create -n ridnet python3.9 conda activate ridnet pip install torch1.13.1cu117 torchvision0.14.1cu117 --extra-index-url https://download.pytorch.org/whl/cu117 pip install numpy opencv-python scikit-image tqdm再说一下为什么推荐1.13.1。我用PyTorch 2.0以上的版本也测过损失收敛曲线差别不大但如果你要加载之前别人训练好的模型权重可能会遇到weights_only参数变化带来的兼容性警告虽然不影响推理但确实会很烦。版本锁死之后整个训练和评估流程我实测下来非常稳定没有出现算子不兼容或者梯度异常的问题。2.2 CPU环境能跑吗如果你的机器没有NVIDIA显卡CPU模式也能完成推理但训练基本不建议。RIDNet参数量大约在1.5M左右batch size16输入128×128的patch在1080Ti上每个epoch大概5分钟CPU的话可能要按小时计算。推理倒是没问题一张512×512的图像CPU环境下处理时间大概在2到3秒可以接受。如果是纯CPU环境安装PyTorch用CPU版本pip install torch1.13.1别装带cu117的版本装错的话import torch会直接闪退。这个坑我在Windows机器上踩过报错信息还不明显。3. RIDNet网络结构拆解EAM模块与残差学习的实现3.1 网络整体架构RIDNet的流程可以分成三段浅层特征提取、深层特征学习和重建输出。输入一张带噪图像首先经过一个3×3卷积得到64通道的浅层特征。然后送入若干个级联的EAM残差块每个残差块内部是一个3×3卷积、ReLU、3×3卷积的结构外圈加了一条旁路把输入直接加到输出上。特征学习完成后经过一个3×3卷积把通道数降回去最后和输入图像做全局残差相加得到去噪结果。这里最关键的设计是全局残差。RIDNet学习的是噪声图而不是干净图也就是网络内部预测的是输入和干净图像之间的差异。这一招的好处是网络只需要拟合高频噪声分布不需要重建图像的主体结构训练难度大幅下降收敛速度明显更快。3.2 EAM模块原理解读EAM全称Efficient Attention Module属于通道注意力机制的一种。它的设计初衷是在SE Block和CBAM之间取一个折中SE Block太重全连接层会引入大量参数CBAM既有通道注意力又有空间注意力结构稍显复杂。EAM的做法是用1×1卷积替代全连接层对全局平均池化后的特征做通道权重预测然后用Sigmoid生成0到1之间的掩码乘回原始特征。这样的设计在数学上非常直观class EAM(nn.Module): def __init__(self, channels): super(EAM, self).__init__() self.pool nn.AdaptiveAvgPool2d(1) self.conv1 nn.Conv2d(channels, channels // 4, 1, biasFalse) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(channels // 4, channels, 1, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): shortcut x y self.pool(x) y self.conv1(y) y self.relu(y) y self.conv2(y) y self.sigmoid(y) return x * y中间瓶颈层的通道数设为输入的1/4这样既保留了通道间非线性交互能力又不会让参数量膨胀。放在残差块内部的话通常是先过两个3×3卷积再接一个EAM最后残差相加。3.3 为什么EAM比SE更适合去噪SE Block对通道做缩放时会把空间信息压缩成一个标量然后通过两个全连接层重新分配权重。这个逻辑在分类任务里没问题但去噪任务是像素级预测空间细节至关重要过度压缩空间信息会损失边缘纹理。EAM的设计在池化这一步用了AdaptiveAvgPool2d(1)但随后的权重预测用卷积实现梯度回传路径更平滑同时1×1卷积天然支持任意空间尺寸输入部署时不需要固定输入大小。我做过对比实验把EAM换成SE Block后SIDD验证集上的PSNR掉了0.2dB左右。原因是SE Block的全连接层会强制引入局部连接假设而卷积形式的通道注意力更符合图像处理的平移等变性。3.4 完整模型代码class RIDNet(nn.Module): def __init__(self, in_channels3, num_features64, num_blocks4): super(RIDNet, self).__init__() self.conv_head nn.Conv2d(in_channels, num_features, 3, padding1) self.body nn.Sequential(*[ ResidualBlockWithEAM(num_features) for _ in range(num_blocks) ]) self.conv_tail nn.Conv2d(num_features, in_channels, 3, padding1) self.relu nn.ReLU(inplaceTrue) def forward(self, x): res x x self.relu(self.conv_head(x)) x self.body(x) x self.conv_tail(x) return x resnum_blocks我默认设成4这是原始论文里用的配置。如果你想让模型更强可以调到6或8代价是显存占用和推理时间线性上升。我用4个块在1080Ti上batch size16显存占用约7GB属于非常友好的范围。4. PSNR/SSIM评估指标原理与代码实现4.1 PSNR的计算细节PSNR峰值信噪比是最常用的全参考图像质量指标单位是dB数值越大代表重建图像越接近参考图像。它基于像素级均方误差PSNR 10 * log10(MAX^2 / MSE)其中MAX是像素值的最大可能取值8位图像取255。如果图像是浮点型数据且范围在0到1之间MAX应取1。我在计算PSNR时踩过一个坑用opencv读图默认是BGR通道顺序用PIL读图是RGB如果评估代码里没统一通道顺序PSNR值会差一点虽然数值上可能只差0.01dB但做论文实验时这种误差是不能接受的。建议在加载数据集时就固定使用一个库全程保持一致。还有一个更重要的细节评估时通常要在Y通道亮度通道上计算PSNR而不是RGB三通道直接平均。原因在于人眼对亮度变化更敏感学术界评估去噪效果时标准做法是把RGB转成YCbCr只取Y通道计算PSNR和SSIM。def rgb_to_ycbcr(img): # img: numpy array, shape [H, W, 3], dtypeuint8, range [0, 255] r, g, b img[:, :, 0], img[:, :, 1], img[:, :, 2] y 0.299 * r 0.587 * g 0.114 * b cb 128 - 0.168736 * r - 0.331264 * g 0.5 * b cr 128 0.5 * r - 0.418688 * g - 0.081312 * b return y, cb, cr def calculate_psnr(img1, img2, border0): if border 0: img1 img1[border:-border, border:-border] img2 img2[border:-border, border:-border] mse np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2) if mse 0: return float(inf) return 10. * np.log10(255. * 255. / mse)4.2 SSIM的计算细节SSIM结构相似性从亮度、对比度和结构三个维度衡量两幅图像的相似度取值范围在-1到1之间越接近1表示越相似。它用局部窗口通常是高斯加权窗口统计均值、方差和协方差再加常数项防止分母为零。实现方式有两种一种是用scikit-image的structural_similarity函数一行代码搞定另一种是自己实现滑动窗口计算。scikit-image版本需要注意参数变化新版API中参数名从win_size改成了win_size兼容data_range必须显式指定否则会报警告。from skimage.metrics import structural_similarity as ssim def calculate_ssim(img1, img2, border0): if border 0: img1 img1[border:-border, border:-border] img2 img2[border:-border, border:-border] return ssim(img1, img2, data_range255)如果你不想依赖scikit-image也可以自己实现核心就是高斯核和局部统计量。我建议在生产环境直接用scikit-image它内部做了很多数值稳定性处理比自己写滑动窗口更可靠。4.3 边界裁剪问题这是一个非常容易混淆的地方。RIDNet的卷积层都用了padding1网络本身不会改变图像尺寸但有些论文在评估时会裁剪掉14个像素的边界因为推理时感受野边缘的像素重建质量较差。我在复现时发现是否裁剪边界对PSNR的影响大约在0.05到0.1dB之间。如果你要和论文报告的数值对比一定要确认对方是否裁剪了边界、裁剪了多少像素。最稳妥的做法是在评估代码里加一个border参数默认设为0需要时手动指定。4.4 评估脚本的完整流程我工程的evaluate.py做的事情是读取干净图像和去噪图像转成Y通道计算PSNR和SSIM打印结果并保存到一个文本文件里。对大尺寸图像还会切成多个patch分别计算然后取平均避免一次性计算过大的mse导致精度问题。5. 从零训练数据准备、训练配置与优化细节5.1 数据集选择与预处理RIDNet原始论文主要在SIDD数据集上做真实图像去噪训练因为SIDD包含真实传感器噪声的配对图像不是合成噪声。SIDD训练集中有大约20000张带噪图像但原始数据是巨大的mat文件直接用会非常吃内存。我的做法是先把SIDD的数据转成png格式的图对noisy和gt各一张然后按场景划分训练集和验证集这样数据加载速度更快也更方便调试。如果暂时拿不到SIDD数据集可以用BSD400配合高斯噪声做合成噪声训练模型结构不受影响但泛化到真实图像上的效果会有差距。我当时是先跑通BSD400的流程之后才在SIDD上微调。这个渐进式的方案对复现来说成本最低。数据增强方面我用的是随机旋转和水平翻转叠加随机裁剪。每次迭代时从大图中随机裁剪128×128的patch同时把长边缩放到256~512之间再裁变相扩大了数据多样性。训练时每个epoch重新生成裁切位置相当于在无限增广的样本上训练。5.2 训练配置参数一览我采用的训练配置如下表所示这套参数在1080Ti上单卡可以稳定复现出论文量级的效果配置项数值说明输入patch大小128×128太大显存吃紧太小感受野不够batch size16受显存16GB限制再大梯度更稳优化器Adambetas(0.9, 0.999)初始学习率1e-3预热2000步后达到学习率策略余弦退火最小降到1e-5损失函数L1 Loss比L2收敛更稳训练轮数100每轮结束后跑验证集梯度裁剪最大范数0.5防止偶然大梯度导致发散损失函数选L1而不是L2是我反复实验后的结论。L2损失在训练初期梯度太大容易把权重撑爆虽然L2理论上在正态噪声假设下是最大似然估计但真实图像噪声并非纯高斯分布L1在实践中的鲁棒性更好。我在SIDD上切换损失函数后PSNR反而涨了0.1dB左右。5.3 训练循环的关键代码for epoch in range(start_epoch, total_epochs): model.train() train_loss 0 for noisy, clean in train_loader: noisy noisy.cuda() clean clean.cuda() output model(noisy) loss criterion(output, clean) optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm0.5) optimizer.step() train_loss loss.item() scheduler.step() if epoch % 10 0 or epoch total_epochs - 1: save_checkpoint(model, epoch, loss)训练日志记录我建议定期输出三个指标当前epoch的平均loss、学习率、验证集PSNR。很多复现失败的项目问题不是模型结构而是可视化不足导致模型发散时根本发现不了。5.4 模型文件的保存与加载训练好的模型权重我统一存成.pth文件里面只保存state_dict不保存整个模型对象。这样文件体积小加载时也更安全。torch.save({ epoch: epoch, state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_psnr: best_psnr, }, checkpoints/ridnet_sidd.pth)加载时注意一点如果你用模型文件做推理直接torch.load后取state_dict再用load_state_dict加载。如果你用的是PyTorch 2.6以上版本torch.load默认的weights_only参数发生了变化可能导致包含优化器状态的checkpoint加载失败这种情况需要显式设置weights_onlyFalse。这也是我推荐用1.13.1的原因之一省心。6. 真实图像去噪模型加载与推理实战6.1 推理demo流程推理阶段的核心就是把训练好的模型权重加载进来对输入图像做归一化、前向传播、反归一化然后保存结果。我写了一个denoise_demo.py支持命令行传入图像路径python denoise_demo.py --input noisy.png --output clean.png --checkpoint checkpoints/ridnet_sidd.pth --cuda程序内部流程如下读取图像转为RGB且数值范围在0到1之间的float32张量。扩batch维度形状从[H, W, 3]变成[1, 3, H, W]。送入模型得到输出。裁剪掉多余维度转回numpy数值范围还原到0~255保存。这里有一个值得强调的操作输入模型前图像要除以255归一化到[0,1]推理完再乘以255。有些复现代码会忘记这一步直接用0~255的原始值输入导致模型输出全是NaN因为网络内部权重是按[0,1]输入训练的。6.2 大图处理方案如果输入图像是4000×3000这种超大尺寸直接把整张图送进网络显存很容易爆掉。我的做法是用重叠patch推理把大图切成512×512的小块每个块和周围重叠32个像素推理完成后只取每块中间区域拼回完整图像。重叠区域取中间的做法能有效避免拼接缝。def denoise_large_image(model, img, patch_size512, overlap32): h, w, c img.shape stride patch_size - overlap output np.zeros_like(img) count np.zeros((h, w, 1)) for i in range(0, h - patch_size 1, stride): for j in range(0, w - patch_size 1, stride): patch img[i:ipatch_size, j:jpatch_size] out_patch model_inference(model, patch) output[i:ipatch_size, j:jpatch_size] out_patch count[i:ipatch_size, j:jpatch_size] 1 output output / count return output对于边缘不足patch_size的区域需要对图像做镜像padding推理完再裁掉。这个细节我一开始没注意导致图像四边出现黑色条纹排查了好久才意识到是padding类型的问题。6.3 真实图像去噪效果观察用训练好的模型跑真实室内照片能明显感觉到三个方面改善高ISO下的彩色噪声被压得很干净暗部区域的色斑几乎消失边缘细节没有出现明显的模糊或振铃。需要注意RIDNet的强项是真实传感器噪声如果你拿它去处理高斯噪声严重合成的图效果不会比DnCNN好太多因为它内部的特征分布更贴近真实噪声源。如果你手头只有一张带噪图没有干净参考图PSNR和SSIM算不出来只能主观判断去噪效果。这时候可以把去噪前后图像拼在一起对比或者用放大局部区域的方式观察纹理保留程度。7. 复现过程中的坑与排查记录7.1 训练Loss不下降这是我复现时遇到的第一个大坑。模型结构没问题数据加载也没问题但训练到第5个epoch时loss就开始震荡不停PSNR在24dB左右死活上不去。排查下来发现是学习率太大了Adam虽然自适应调整但初始学习率1e-3配合L1损失在SIDD这种真实噪声数据上还是偏激进。换成预热策略后前2000步用线性从1e-4升到1e-3之后再用余弦退火问题直接解决。7.2 验证集PSNR和论文对不上论文报告SIDD上的PSNR在39dB左右我一开始只能跑到36.5dB差了一大截。检查后发现问题出在评估预处理上论文在评估SIDD时并不是直接在原始分辨率评估而是先对图像做双三次下采样再评估。另外我评估时没有裁剪边界而论文可能用了不同评估代码。最后统一了数据预处理流程和边界裁剪规则PSNR才对齐到38.9dB左右。7.3 模型加载时权重维度不匹配训练中途改了网络结构的通道数之后加载旧权重就一直报size mismatch。这个问题完全是自己粗心模型定义和权重没对应起来。建议每次修改网络结构后重新训练或使用官方权重不要混用。如果你加载别人给的模型文件时遇到这个报错大概率是对方用的网络配置和你定义的不一致检查一下EAM模块的通道缩减比例我代码里是1/4。7.4 显存不足训练时batch size16输入128×128在1080Ti上没问题。但如果换到6GB显存的显卡可以把batch size降到8patch_size降到96同时开启torch.backends.cudnn.benchmarkTrue能节省约30%显存。对最终模型效果影响很小因为去噪网络对patch大小不敏感。7.5 常见问题速查表现象原因解决方案输出全是黑色忘了反归一化推理结果乘以255输出全是NaN输入未归一化输入除以255训练Loss震荡学习率过大使用预热余弦退火四边有黑边padding类型错误使用reflect padding推理后裁剪权重加载报错版本或结构不匹配统一网络配置或设置weights_onlyFalse拼接缝明显大图切块没有重叠推理加入overlap参数并取中间区域7.6 关于模型文件的几点使用建议附带的训练好权重是我在SIDD数据集上训练100个epoch得到的可以直接用于真实图像去噪推理。如果你想在它基础上做二分类任务微调或者接入自己的项目建议先跑一遍demo确认环境正常再改输入输出格式。这个权重的license仅用于学术研究如果要在商业场景使用还是建议在合规的数据集上重新训练。最后再说一点个人经验复现论文模型很多时候难点不在网络结构而在数据、训练细节、评估方式这些容易被忽略的环节。我在这次RIDNet复现过程中最深的体会是评估代码一定要先写对再谈训练。如果你算PSNR和SSIM本身的实现就有偏差模型训得再好指标也是假的。我的建议是先写一段评估脚本拿一幅已知图像做输入确认输出和你手算的数值一致再开始跑训练。另外工程里的模型文件、训练脚本和推理demo都是我实际跑通过的版本不是摆样子的空架子。你拿到之后可以先跑demo看效果再根据自己需求调整网络宽度和深度。RIDNet的结构非常简单改起来也很顺手非常适合作为真实图像去噪方向的入门baseline使用。如果你后面想做得更深可以试试在这个基础上加多尺度结构或者拿它做视频去噪的backbone都有不少扩展空间。本文还有配套的精品资源点击获取