ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

DnCNN复现指南:从DnCNN-B到DnCNN-3的四个变体实战

DnCNN复现指南:从DnCNN-B到DnCNN-3的四个变体实战 简介本资源面向图像去噪方向的深度学习学习者与研究者提供基于PyTorch的DnCNN完整复现代码并在原始DnCNN基础上扩展实现了DnCNN-B、CDnCNN-B与DnCNN-3的训练与测试流程适合具备一定PyTorch基础、希望系统复现论文实验的读者。压缩包共20个文件约13.2MB包含6个pth模型权重、5个py脚本、4个xml配置及md说明等覆盖数据制作、模型训练、测试评估与结果保存等环节并附带训练好的权重文件模型性能与论文近似可不训练直接测试。资源几乎实现论文中全部图表相当于把整套工作完整做了一遍读者可据此更换路径与参数训练自己的图像数据集并通过测试脚本计算平均PSNR与SSIM。目前已有2304人学习下载适合作为图像去噪入门复现与实验对照的参考。1. DnCNN 复现这件事从一份源码到四个变体值不值得动手如果你手上只有一份 DnCNN 的 PyTorch 源码跑通一次训练和测试其实半天就能搞定。但真正让这个项目值得投入的是标题里那串后缀DnCNN-B、CDnCNN-B、DnCNN-3。它们不是简单的改名而是原作者在论文里针对不同噪声假设、不同盲去噪场景、不同任务泛化能力做的三组关键实验。很多人复现 DnCNN 只跑了灰度图加固定高斯噪声然后发现 PSNR 对不上论文就开始怀疑人生。问题往往不在代码而在于没搞清楚 DnCNN-B 的盲噪声训练策略、CDnCNN-B 的颜色通道处理方式以及 DnCNN-3 是怎么用一个模型同时干三件事的。这篇笔记面向的是已经装好 PyTorch 环境、想把这四个变体一次性复现出来的从业者。我会按「先讲清每个变体在论文里到底改了什么再给可抄的代码和参数最后把训练中真正会翻车的地方摊开说」的顺序来写。读完你应该能自己搭出一套完整的训练测试流水线而不是只跑通一个 demo。PyTorch 环境搭建、CUDA 版本匹配这些基础操作网上教程很多这里不展开默认你已经能跑torch.cuda.is_available()返回 True。2. 四个变体到底差在哪DnCNN、DnCNN-B、CDnCNN-B、DnCNN-3 的设计动机2.1 从 DnCNN 到 DnCNN-B噪声水平图是怎么塞进网络的原始 DnCNN 针对的是非盲去噪也就是训练时噪声水平 σ 是已知的网络通过一个噪声水平图noise level map把 σ 作为额外输入告诉模型。具体做法是把 σ 归一化后复制成与输入图像同尺寸的单通道图和灰度图拼接成两通道输入。这个设计在 σ 固定时效果很好但一旦测试噪声水平和训练不一致性能就掉得厉害。DnCNN-B 的 B 代表 Blind核心改动是训练时 σ 在一个范围内随机采样比如 [0, 55]而不是固定值。这样网络见过了各种噪声强度测试时不需要再喂噪声水平图输入退化成单通道灰度图。代价是同一噪声水平下的峰值性能会略低于非盲版本但泛化能力大幅提升。复现时最容易搞错的一点是DnCNN-B 的输入通道数要改成 1而原始 DnCNN 是 2。这个改动会连锁影响第一层卷积的in_channels和后续所有和通道数相关的配置。2.2 CDnCNN-B 的颜色通道处理为什么不是简单把三通道拼起来CDnCNN-B 的 C 代表 Color。最直觉的做法是把 RGB 三通道直接作为输入输出也是三通道。但论文里的做法更细它仍然只预测噪声残差但训练时对每个颜色通道独立加噪或者对三通道联合加噪。实际复现时我一般会采用三通道输入三通道输出的结构第一层卷积in_channels3最后一层out_channels3。关键区别在于数据加载阶段要对 RGB 图像的三个通道分别施加高斯噪声而不是只对灰度通道加。这里有个容易忽略的细节CDnCNN-B 的盲训练策略和 DnCNN-B 一致σ 在 [0, 55] 随机采样。但颜色图像的噪声方差通常比灰度图更复杂因为 RGB 通道间可能存在相关性。如果你的数据集是自然图像建议先用cv2.cvtColor确认通道顺序是 RGB 而不是 BGR否则训练出来的模型颜色会偏。这个坑我在第一次复现时踩过PSNR 看着还行但输出图像整体偏蓝排查了半天才发现是 OpenCV 默认 BGR 导致的。2.3 DnCNN-3 的多任务设计一个模型干三件事的代价DnCNN-3 的 3 指的是三个任务灰度去噪、颜色去噪、单图像超分辨率。论文里的做法是训练一个模型同时处理这三类任务通过不同的数据预处理和损失函数来区分。复现时最直接的方式是构造一个多任务数据加载器每个 batch 里混合三种任务的数据但这样训练会很不稳定。我一般会采用分阶段训练先单独训练去噪任务再逐步加入超分数据最后联合微调。DnCNN-3 的网络结构和 DnCNN-B 基本一致区别在于训练数据的组织和损失函数的加权。超分任务需要把低分辨率图像上采样到高分辨率再让网络学习残差。这里要注意超分任务的输入尺寸和去噪任务不同如果混在一个 batch 里需要统一 resize 到相同尺寸否则会报维度错误。我通常会把所有输入统一到 40x40 的 patch这样既能覆盖去噪也能覆盖 2 倍超分的感受野需求。3. 把源码跑起来数据准备、模型定义和训练循环的最小实现3.1 数据加载怎么构造带噪声的 patch 对DnCNN 系列训练时用的是 patch 级别的数据而不是整图。常见做法是从训练集图像中随机裁剪 40x40 或 50x50 的 patch然后在线加噪。这样既能扩充数据量又能控制显存占用。下面是一个可复用的 Dataset 实现支持灰度、颜色和盲噪声三种模式。import torch from torch.utils.data import Dataset import cv2 import numpy as np import random class DnCNNDataset(Dataset): def __init__(self, image_paths, patch_size40, modegray, sigma_range(0, 55)): mode: gray 对应 DnCNN/DnCNN-B, color 对应 CDnCNN-B sigma_range: 盲训练时的噪声水平范围非盲时传固定值如 (25, 25) self.image_paths image_paths self.patch_size patch_size self.mode mode self.sigma_range sigma_range def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img cv2.imread(self.image_paths[idx]) if self.mode gray: img cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) else: img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w img.shape[:2] # 随机裁剪 patch如果图像太小就 resize if h self.patch_size or w self.patch_size: img cv2.resize(img, (max(h, self.patch_size), max(w, self.patch_size))) h, w img.shape[:2] top random.randint(0, h - self.patch_size) left random.randint(0, w - self.patch_size) patch img[top:topself.patch_size, left:leftself.patch_size].astype(np.float32) / 255.0 # 随机采样噪声水平 sigma random.uniform(*self.sigma_range) / 255.0 noise np.random.randn(*patch.shape) * sigma noisy patch noise # 转为 tensor灰度图增加通道维 if self.mode gray: clean_t torch.from_numpy(patch).unsqueeze(0) noisy_t torch.from_numpy(noisy).unsqueeze(0) else: clean_t torch.from_numpy(patch).permute(2, 0, 1) noisy_t torch.from_numpy(noisy).permute(2, 0, 1) return noisy_t, clean_t这段代码的关键参数有三个patch_size决定感受野和显存占用40 是论文里的常用值sigma_range在盲训练时设为 (0, 55)非盲时设为固定值mode控制灰度还是颜色。注意噪声是在 [0,1] 归一化后加的σ 也要除以 255否则量级不对训练会直接发散。这个细节很多复现代码里写错了导致 loss 一直降不下去。3.2 模型定义17 层卷积的通道数和残差连接DnCNN 的结构很规整第一层 ConvReLU中间 15 层 ConvBNReLU最后一层 Conv 输出残差。下面是一个支持灰度、颜色和不同深度的实现。import torch.nn as nn class DnCNN(nn.Module): def __init__(self, channels1, num_layers17, features64): channels: 1 对应灰度3 对应颜色 num_layers: 论文默认 17 features: 中间层通道数默认 64 super(DnCNN, self).__init__() layers [] # 第一层Conv ReLU layers.append(nn.Conv2d(channels, features, kernel_size3, padding1, biasFalse)) layers.append(nn.ReLU(inplaceTrue)) # 中间层Conv BN ReLU for _ in range(num_layers - 2): layers.append(nn.Conv2d(features, features, kernel_size3, padding1, biasFalse)) layers.append(nn.BatchNorm2d(features)) layers.append(nn.ReLU(inplaceTrue)) # 最后一层Conv输出残差 layers.append(nn.Conv2d(features, channels, kernel_size3, padding1, biasFalse)) self.dncnn nn.Sequential(*layers) def forward(self, x): # 输入是噪声图输出是预测的噪声残差 noise self.dncnn(x) # 干净图 噪声图 - 预测噪声 return x - noise这里有几个参数需要根据变体调整channels在 DnCNN-B 和 DnCNN-3 的灰度任务里是 1在 CDnCNN-B 里是 3num_layers默认 17但 DnCNN-3 有时会用更深的版本features默认 64显存不够可以降到 32但 PSNR 会掉 0.2 左右。注意 forward 里返回的是干净图不是残差这样 loss 可以直接用 MSE 计算。如果你想让模型输出残差把return x - noise改成return noise然后 loss 里手动做减法效果一样但代码可读性差一些。3.3 训练循环损失函数、优化器和学习率调度DnCNN 论文用的是 MSE loss 和 Adam 优化器初始学习率 1e-3每 30 个 epoch 衰减 0.1 倍。下面是一个最小训练循环支持盲训练和非盲训练。import torch from torch.utils.data import DataLoader from torch.optim import Adam from torch.optim.lr_scheduler import MultiStepLR def train_one_epoch(model, dataloader, optimizer, criterion, device): model.train() total_loss 0 for noisy, clean in dataloader: noisy, clean noisy.to(device), clean.to(device) optimizer.zero_grad() output model(noisy) loss criterion(output, clean) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader) # 配置 device torch.device(cuda if torch.cuda.is_available() else cpu) model DnCNN(channels1, num_layers17, features64).to(device) criterion nn.MSELoss() optimizer Adam(model.parameters(), lr1e-3) scheduler MultiStepLR(optimizer, milestones[30, 60, 90], gamma0.1) # 假设 dataset 已经构造好 dataloader DataLoader(dataset, batch_size128, shuffleTrue, num_workers4) for epoch in range(100): loss train_one_epoch(model, dataloader, optimizer, criterion, device) scheduler.step() print(fEpoch {epoch}, Loss: {loss:.6f})batch_size设为 128 是论文里的配置显存不够可以降到 64 或 32但学习率要相应调小。milestones控制学习率衰减的 epoch如果你训练轮数少比如 50 轮可以改成 [20, 40]。num_workers根据 CPU 核数调整一般设为 4 或 8。注意盲训练时每个 batch 里的噪声水平是随机的所以 loss 波动会比非盲训练大这是正常的不要因为 loss 震荡就调小学习率。4. 测试与评估PSNR 怎么算才和论文对得上4.1 PSNR 计算的三个边界条件PSNR 的公式很简单10 * log10(MAX^2 / MSE)但实际算的时候有三个坑。第一MAX 是 1 还是 255取决于你的数据归一化方式如果训练时归一化到 [0,1]测试时也要用 1。第二MSE 是在哪个颜色空间算的灰度图直接在单通道算颜色图要在 RGB 三通道算完取平均而不是转成 YUV 再算。第三边界像素要不要裁剪DnCNN 的卷积有 padding输出尺寸和输入一致不需要裁剪但如果你用了其他结构可能需要裁掉边缘几个像素。import numpy as np def calculate_psnr(img1, img2, max_val1.0): img1, img2: numpy array, shape (H, W) 或 (H, W, C), 范围 [0, max_val] mse np.mean((img1 - img2) ** 2) if mse 0: return float(inf) psnr 10 * np.log10(max_val ** 2 / mse) return psnr这个函数对灰度和颜色都适用只要传入的数组形状一致。注意如果图像是 uint8 类型要先转成 float 并除以 255否则 max_val 要设成 255。我一般会在测试脚本里统一做归一化避免混淆。4.2 测试集上的评估流程测试时不需要加噪再评估而是用固定的测试集对每张图施加已知 σ 的噪声然后去噪再算 PSNR。对于盲模型σ 可以设成训练范围外的值比如 60 或 70看泛化能力。下面是一个完整的评估脚本。def evaluate(model, test_paths, sigma25, modegray, devicecuda): model.eval() psnr_list [] with torch.no_grad(): for path in test_paths: img cv2.imread(path) if mode gray: img cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) else: img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.astype(np.float32) / 255.0 noise np.random.randn(*img.shape) * (sigma / 255.0) noisy img noise if mode gray: noisy_t torch.from_numpy(noisy).unsqueeze(0).unsqueeze(0).to(device) else: noisy_t torch.from_numpy(noisy).permute(2, 0, 1).unsqueeze(0).to(device) output model(noisy_t) output output.squeeze().cpu().numpy() if mode color: output output.transpose(1, 2, 0) psnr calculate_psnr(img, output) psnr_list.append(psnr) return np.mean(psnr_list)这个脚本里sigma是测试噪声水平盲模型可以设成任意值。注意noisy_t的维度灰度图是 (1, 1, H, W)颜色图是 (1, 3, H, W)。如果报维度错误先检查这里。另外测试时不要用model.train()否则 BatchNorm 会更新 running stats导致结果不稳定。5. 避坑与排查复现 DnCNN 系列时最容易翻车的五个地方5.1 现象训练 loss 降到 0.01 以下但测试 PSNR 只有 20 出头原因通常是数据归一化不一致。训练时如果用了ToTensor()把图像转到 [0,1]但测试时忘了除 255或者反过来模型学到的残差量级就完全错了。解决方法是统一在 Dataset 里做归一化测试脚本里也显式除以 255不要依赖任何隐式转换。5.2 现象CDnCNN-B 训练时 loss 正常下降但输出图像颜色失真这个坑我踩过两次。第一次是 OpenCV 读图默认 BGR训练时没转 RGB模型学到的通道顺序和测试时不一致。第二次是 BatchNorm 在颜色通道上独立计算如果 batch 里颜色分布不均衡running stats 会偏。解决方法是在 Dataset 里强制cv2.cvtColor(img, cv2.COLOR_BGR2RGB)并且训练时 shuffle 要开batch_size 不要太小至少 64。5.3 现象DnCNN-3 多任务训练时去噪性能突然掉点原因是超分任务的输入尺寸和去噪任务不同如果混在一个 batch 里resize 会引入额外插值噪声干扰去噪任务的学习。我一般会分阶段训练前 50 轮只训去噪中间 30 轮加入超分但降低超分 loss 权重最后 20 轮联合微调。另外超分任务的低分辨率图像要用双三次插值下采样不要用最近邻否则边缘伪影会被模型当成噪声学进去。5.4 现象盲训练时 σ 范围设成 [0, 55]但测试 σ50 时 PSNR 比 σ25 还低这是正常的因为盲模型在中间噪声水平上性能最稳两端会掉。如果你希望 σ50 时性能更好可以把训练范围改成 [0, 70]但 σ10 时的性能会略降。这是一个权衡没有免费午餐。我一般会根据实际应用场景选范围如果主要处理中等噪声[0, 55] 就够了。5.5 现象训练到一半 loss 突然变成 NaN最常见的原因是学习率太大或者噪声水平采样到了极端值导致梯度爆炸。解决方法是在train_one_epoch里加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。另外检查sigma_range的上限是否超过了 100超过 100 时噪声图几乎全是噪声MSE loss 会很大容易炸。如果必须处理极高噪声先把学习率降到 1e-4。6. 进阶技巧用 ONNX 导出和混合精度训练把复现效率提上去6.1 导出 ONNX 时怎么处理动态输入尺寸DnCNN 是全卷积网络理论上支持任意尺寸输入但导出 ONNX 时如果写死尺寸部署时换分辨率就会报错。正确做法是用dynamic_axes把高度和宽度设为动态。import torch.onnx model.eval() dummy_input torch.randn(1, 1, 40, 40).to(device) torch.onnx.export( model, dummy_input, dncnn_b.onnx, input_names[input], output_names[output], dynamic_axes{input: {2: height, 3: width}, output: {2: height, 3: width}}, opset_version11 )导出后可以用onnxruntime验证一下不同尺寸的输入是否能跑通。注意 BatchNorm 在导出时会折叠进卷积所以 ONNX 模型比 PyTorch 模型略小推理速度也更快。如果导出失败先检查opset_version11 以上对动态尺寸支持更好。6.2 混合精度训练省显存但不省精度的几个参数PyTorch 的amp模块可以自动做混合精度但 DnCNN 这种回归任务对数值精度敏感直接用autocast可能会导致 PSNR 掉 0.1 到 0.3。我的做法是只在 forward 和 loss 计算时用autocastbackward 和 optimizer step 保持 FP32。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for noisy, clean in dataloader: noisy, clean noisy.to(device), clean.to(device) optimizer.zero_grad() with autocast(): output model(noisy) loss criterion(output, clean) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这样显存能省 30% 左右batch_size 可以翻倍训练速度提升 20% 到 40%。如果发现 PSNR 掉了超过 0.5就把autocast关掉或者只对卷积层用。另外GradScaler在训练初期会动态调整 scale如果 loss 出现 inf它会自动跳过这个 batch不会导致训练崩溃。6.3 一个我常用的验证习惯每次改完模型结构或训练参数我不会直接跑完整训练集而是先拿 10 张图跑 5 个 epoch看 loss 是否正常下降再用测试脚本算 PSNR。如果 5 个 epoch 后 PSNR 能到 25 以上说明配置基本正确可以开完整训练。这个习惯帮我省了很多次通宵跑完发现参数写错的后悔药。DnCNN 系列复现不难难的是细节一致性希望帮到你。本文还有配套的精品资源点击获取
返回列表