DnCNN图像去噪实战：从DnCNN-B到DnCNN-3的PyTorch复现与避坑指南

简介本资源面向图像去噪方向的深度学习学习者与研究者提供基于PyTorch的DnCNN完整复现代码并在原始DnCNN基础上扩展实现了DnCNN-B、CDnCNN-B与DnCNN-3的训练与测试流程适合具备一定PyTorch基础、希望系统复现论文实验的读者。压缩包共20个文件约13.2MB包含6个pth模型权重、5个py脚本、4个xml配置及md说明等覆盖数据生成、训练、测试与结果保存等环节模型权重可直接加载测试无需重新训练。目前已有2304人学习下载。资源几乎复现了论文中的全部图表相当于把整套工作完整做了一遍读者可据此更换路径与参数训练自己的图像数据集并通过测试脚本计算平均PSNR与SSIM快速验证不同噪声水平下的去噪效果兼具参考价值与实操性。1. 从一份能跑通的 DnCNN 复现代码说起它到底补了哪些坑如果你搜过 DnCNN 的 PyTorch 实现大概率见过那种「只有模型定义、没有训练脚本、测试靠脑补」的仓库。这份资源不一样它是在一份能跑通的 DnCNN 源码基础上把 DnCNN-B、CDnCNN-B、DnCNN-3 三个变体的训练和测试流程都补齐了。换句话说你拿到的不只是一个model.py而是一套从数据准备、噪声合成、训练调度到 PSNR 评估的完整链路。适合谁正在做图像去噪 baseline 对比的研究生、需要快速验证去噪网络结构的算法工程师以及想拿 DnCNN 当 backbone 改注意力模块的人。它解决的核心问题是让你不用从零搭训练框架直接改配置就能复现论文里的盲去噪和非盲去噪结果。下面我按「先搞懂变体差异再动手跑通最后避开几个血泪坑」的顺序拆一遍。2. 四个变体到底差在哪DnCNN、DnCNN-B、CDnCNN-B、DnCNN-3 的选型逻辑2.1 从 DnCNN 到 DnCNN-B噪声水平输入方式的改变原始 DnCNN 是非盲去噪意思是训练和测试时都要把噪声标准差 σ 作为先验告诉网络。具体做法是把 σ 归一化后拼成一个常数通道和噪声图像一起送进网络。这份源码里对应的是DnCNN类输入通道数为 2灰度图 1 通道 σ 图 1 通道。而 DnCNN-B 是盲去噪版本网络不知道 σ 是多少输入通道数变成 1靠网络自己从图像里估计噪声水平。这个改动看似只是少了一个通道但训练策略完全不同DnCNN 通常针对单个 σ 训练一个模型DnCNN-B 则要在一定范围的 σ 上混合训练才能覆盖未知噪声。选型建议很直接如果你的测试集噪声水平已知且固定用 DnCNN 非盲版本收敛更快、PSNR 更高如果噪声水平未知或变化必须上 DnCNN-B。我一般会先跑 DnCNN-B 看盲去噪的 baseline再针对特定 σ 微调非盲版本做上限对比。2.2 CDnCNN-B 与 DnCNN-3彩色图像和三种噪声水平的处理CDnCNN-B 里的 C 就是 Color把输入从单通道灰度扩展到三通道 RGB。这里有个容易翻车的点彩色图像的噪声通常在 RGB 三个通道独立添加但网络是否共享权重、是否在 YCbCr 空间处理会直接影响结果。这份源码采用的是直接在 RGB 空间训练卷积核在三个通道上共享参数量和不共享相比少很多适合显存有限的场景。DnCNN-3 则是「一个模型打三种噪声」高斯去噪、SIDD 真实去噪、JPEG 去块效应。它的网络结构和 DnCNN-B 一致区别在训练数据混合了三种退化类型。如果你只想快速验证一个模型能否同时处理多种退化DnCNN-3 是最省事的起点。但要注意混合训练时三种任务的损失量级可能不一致源码里默认用 MSE实际跑的时候可能需要给不同任务加权。2.3 参数量与感受野的取舍DnCNN 系列都是 17 层卷积前 16 层是 ConvBNReLU最后一层是单 Conv 输出残差。感受野大约 35×35对于 40×40 的 patch 刚好覆盖。这意味着训练时 patch size 不能太小否则边缘信息学不到也不能太大否则显存吃紧。源码默认 patch size 是 40batch size 16这个配置在 8GB 显存的卡上跑 DnCNN-B 没问题但 CDnCNN-B 因为三通道输入显存占用大约翻三倍需要把 batch size 降到 4 或 8。提示如果你打算把 DnCNN 当 backbone 加注意力模块建议先冻结前 10 层跑几个 epoch确认梯度不爆炸再解冻全量微调。3. 把代码跑起来环境、数据合成与训练脚本的实操步骤3.1 环境搭建与依赖版本这份源码基于 PyTorch没有用复杂的第三方框架。我实测过的稳定组合是 Python 3.8 PyTorch 1.12 CUDA 11.3或者 PyTorch 2.0 CUDA 11.8。如果你在 WSL 里配环境注意 CUDA 驱动版本要和 Windows 主机一致否则torch.cuda.is_available()会返回 False。安装命令如下# 创建虚拟环境避免和系统包冲突 conda create -n dncnn python3.8 -y conda activate dncnn # 安装 PyTorch以 CUDA 11.3 为例 pip install torch1.12.1cu113 torchvision0.13.1cu113 -f https://download.pytorch.org/whl/torch_stable.html # 安装其他依赖 pip install numpy opencv-python scikit-image tqdm tensorboard逻辑说明torch和torchvision版本必须匹配cu113表示 CUDA 11.3 编译版本。如果你用 7900XTX 这类 AMD 卡PyTorch 官方不支持 ROCm 以外的加速建议直接用 CPU 跑小规模验证或者换 N 卡。scikit-image用来算 PSNR 和 SSIMtensorboard看训练曲线。3.2 训练数据合成高斯噪声怎么加才和论文一致DnCNN 原文用的是 BSD400 Waterloo 数据集合成高斯噪声的方式是img sigma/255 * randn。注意这里 sigma 是 0-255 尺度的不是 0-1。很多复现翻车就翻在这里如果先把图像归一化到 [0,1] 再加噪声sigma 要除以 255。源码里提供了data_generator.py核心逻辑如下import numpy as np import cv2 def add_gaussian_noise(img, sigma): img: 灰度或彩色图像uint8范围 0-255 sigma: 噪声标准差0-255 尺度 返回加噪后的 uint8 图像 noise np.random.randn(*img.shape) * sigma noisy img.astype(np.float32) noise # 裁剪到合法范围避免溢出 noisy np.clip(noisy, 0, 255).astype(np.uint8) return noisy # 生成训练 patch 的示例 def generate_patches(img, patch_size40, stride10, sigma25): h, w img.shape[:2] patches [] 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] noisy_patch add_gaussian_noise(patch, sigma) patches.append((noisy_patch, patch)) return patches参数说明sigma25对应论文里常用的噪声水平patch_size40是 DnCNN 的标准输入stride10控制 patch 重叠程度stride 越小样本越多但训练越慢。注意np.clip这一步不能省否则加噪后像素值可能超过 255 或低于 0训练时 loss 会异常。3.3 训练脚本的关键参数与启动命令源码里训练入口是train.py支持通过 argparse 切换模型类型。启动 DnCNN-B 盲去噪训练的命令python train.py \ --model dncnn_b \ --data_dir ./data/BSD400 \ --sigma_min 10 \ --sigma_max 50 \ --patch_size 40 \ --batch_size 16 \ --epochs 50 \ --lr 1e-3 \ --log_dir ./logs/dncnn_b逻辑说明--model指定变体可选dncnn、dncnn_b、cdncnn_b、dncnn_3--sigma_min和--sigma_max只在盲去噪时生效每个 batch 随机采样一个 sigma--lr 1e-3是初始学习率源码默认用 StepLR每 10 个 epoch 乘 0.5。如果你跑 CDnCNN-B把--model换成cdncnn_b同时把--batch_size降到 8否则容易 OOM。训练过程中用 TensorBoard 看 loss 曲线tensorboard --logdir ./logs --port 6006正常情况 loss 在前 5 个 epoch 下降很快之后趋于平缓。如果 loss 震荡严重检查学习率是否太大或者数据里有没有全黑/全白的异常 patch。3.4 测试与 PSNR 评估测试脚本test.py会加载 checkpoint在 Set12 或 BSD68 上计算平均 PSNR。命令示例python test.py \ --model dncnn_b \ --checkpoint ./logs/dncnn_b/model_best.pth \ --test_dir ./data/Set12 \ --sigma 25注意盲去噪模型测试时也要指定--sigma但这个 sigma 只用于合成测试噪声不输入网络。评估指标除了 PSNR源码还输出了 SSIM方便和论文表格对齐。如果你发现 PSNR 比论文低 0.5dB 以上优先检查测试图像是否做了边界裁剪——DnCNN 输出的是残差边缘几圈像素可能不可靠常见做法是裁掉 17 像素再算指标。4. 避坑与排查复现 DnCNN 时最容易翻车的五个地方4.1 现象训练 loss 正常下降但测试 PSNR 只有 20dB 左右原因数据归一化不一致。训练时图像是 0-255 尺度测试时如果用了ToTensor()自动归一化到 0-1网络看到的输入分布完全不同。解决统一用 0-255 尺度或者在训练和测试里都做相同的归一化并在算 PSNR 前反归一化。4.2 现象CDnCNN-B 训练时显存溢出batch size 降到 1 还是 OOM原因彩色图像三通道输入中间特征图通道数 64显存占用是灰度的三倍。如果输入 patch 是 40×40单张图占用不大但源码里可能默认用了 256×256 的整图训练。解决确认--patch_size是 40 而不是 256同时把--batch_size设为 4 或 8。如果还不行把模型中间通道数从 64 降到 32但 PSNR 会掉 0.2dB 左右。4.3 现象DnCNN-3 在 JPEG 去块任务上 PSNR 远低于论文原因JPEG 压缩质量因子没有对齐。论文里 JPEG 去块用的是 quality factor 10-40 的压缩图像如果你用默认的 75 质量因子退化程度太轻网络学不到东西。解决在数据合成阶段显式指定cv2.imwrite的JPEG_QUALITY参数范围 10-40并且训练时混合多个质量因子。4.4 现象盲去噪模型在 sigma50 时 PSNR 骤降原因训练时 sigma 采样范围太窄。如果--sigma_max只设到 30网络没见过高噪声测试 sigma50 自然崩。解决把--sigma_min和--sigma_max设成 0 和 55覆盖常见噪声范围。但注意 sigma0 时噪声图像等于原图loss 会接近 0可以设成 1 避免除零。4.5 现象训练到一半 loss 变成 NaN原因学习率太大或者梯度爆炸。DnCNN 用了 BN 层理论上对学习率不敏感但如果你的数据里有异常值比如全 255 的 patchBN 的 running mean 会被带偏。解决加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)同时检查数据里有没有坏图。另外把学习率从 1e-3 降到 1e-4 再试。5. 进阶技巧用 DnCNN 做迁移学习和 ONNX 导出5.1 冻结部分层做小样本微调如果你手头只有几十张特定场景的噪声图像从头训练 DnCNN 肯定过拟合。常见做法是加载预训练权重冻结前 10 层只微调后 7 层。代码片段import torch from model import DnCNN # 加载预训练模型 model DnCNN(in_channels1, out_channels1, num_layers17) checkpoint torch.load(./logs/dncnn_b/model_best.pth) model.load_state_dict(checkpoint[state_dict]) # 冻结前 10 层 for i, layer in enumerate(model.features): if i 10: for param in layer.parameters(): param.requires_grad False # 只优化剩余参数 optimizer torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-4)逻辑说明model.features是源码里定义的卷积层序列前 10 层冻结后反向传播只更新后 7 层。学习率用 1e-4 而不是 1e-3避免破坏预训练特征。微调 20 个 epoch 通常就能看到明显提升。5.2 导出 ONNX 做推理加速PyTorch 模型部署到 C 或 TensorRT 时ONNX 是中间格式。导出时注意输入尺寸要固定DnCNN 支持动态尺寸但 ONNX 对动态轴支持有限。导出命令import torch from model import DnCNN model DnCNN(in_channels1, out_channels1, num_layers17) model.load_state_dict(torch.load(./logs/dncnn_b/model_best.pth)[state_dict]) model.eval() # 固定输入尺寸 1x1x256x256 dummy_input torch.randn(1, 1, 256, 256) torch.onnx.export( model, dummy_input, dncnn_b.onnx, input_names[input], output_names[output], opset_version11, dynamic_axes{input: {2: height, 3: width}, output: {2: height, 3: width}} )参数说明opset_version11兼容大多数推理引擎dynamic_axes允许输入高度和宽度动态变化但部分 TensorRT 版本对动态轴支持不好如果部署到 TensorRT建议固定尺寸。导出后用onnxruntime验证输出和 PyTorch 一致import onnxruntime as ort import numpy as np sess ort.InferenceSession(dncnn_b.onnx) input_name sess.get_inputs()[0].name dummy np.random.randn(1, 1, 256, 256).astype(np.float32) onnx_out sess.run(None, {input_name: dummy})[0] torch_out model(torch.from_numpy(dummy)).detach().numpy() print(最大误差:, np.abs(onnx_out - torch_out).max())如果最大误差在 1e-5 以内说明导出成功。超过 1e-3 就要检查是否有不支持的操作被替换成了近似实现。5.3 一个我踩过的坑BN 层在 ONNX 里的融合DnCNN 的 ConvBNReLU 结构在导出 ONNX 时BN 层默认不会自动融合进 Conv。推理时 BN 会单独算一次增加耗时。我一般会在导出前手动调用torch.onnx.utils.remove_experimental_patterns()或者用torch.quantization.fuse_modules把 ConvBNReLU 融合成一个模块。融合后模型大小不变但推理速度能快 15% 左右。从那以后我每次导出 ONNX 都强制走一遍融合检查不然部署到边缘设备上延迟会莫名其妙高出一截。希望帮到你。本文还有配套的精品资源点击获取