简介FFDNet-pytorch 是一份面向图像去噪任务的 PyTorch 实现资源适合具备一定深度学习基础、希望快速复现或集成 FFDNet 模型的研究者与开发者。它针对传统去噪方法计算复杂、对噪声类型敏感等问题提供端到端可训练的网络结构与配套数据可用于高斯噪声、椒盐噪声等多种场景并延伸至医学影像、遥感与视频去噪等应用。资源包共 404 个文件约 374.65MB包含 293 个 png 图像样本、42 个 txt 说明、41 个 pth 权重文件、7 个 py 脚本及 h5 数据文件等覆盖训练日志、预训练模型与代码实现便于直接加载或二次训练。目前已有 928 人学习下载。读者可借此获得完整的网络定义、数据预处理、损失计算与优化流程并参考权重与日志快速验证去噪效果在保留细节与去除噪声之间取得平衡。1. ffdnet-pytorch.zip 拆开之后一个能跑通的 FFDNet 降噪工程长什么样你手上如果有一个叫ffdnet-pytorch.zip的压缩包大概率是两种情况之一要么是从某个开源仓库拉下来的 FFDNet 复现工程要么是自己攒了一半、准备在 PyTorch 上把 FFDNet 跑起来的半成品。FFDNet 这个模型本身不复杂——它是在 DnCNN 基础上做的一个「可调噪声水平」的降噪网络核心思路是把噪声水平图noise level map作为额外输入拼到网络里让一个模型能处理不同强度的噪声而不是每个噪声等级训一个模型。这个设计在 2018 年前后算是很实用的工程妥协放到今天做图像预处理、数据集清洗、低质量图像增强依然有它的位置。ffdnet-pytorch.zip这个标题背后读者真正关心的是三件事这个工程能不能在我的环境里跑起来、我能不能拿它处理自己的图、我改哪里能适配我的噪声场景。所以这篇不聊论文里的公式推导只讲怎么把这个 zip 解压之后一步步跑通参数在哪调以及我踩过的那些坑。适合已经装过 PyTorch、能看懂基本训练循环、但没系统跑过 FFDNet 的人。如果你连 PyTorch 环境都还没搭先把环境搞定再回来后面第 2 章会顺带说清楚版本对应关系。2. 环境与依赖从 pytorch 安装到 ffdnet 能 import 的最小闭环2.1 先确认 PyTorch 版本和 CUDA 的对应关系FFDNet 这类老架构对 PyTorch 版本不挑但对你机器上的 CUDA 和显卡驱动挑。我一般不会一上来就装最新版而是先看显卡。如果你用的是 7900XTX 这类 AMD 卡走 WSL 里跑 PyTorch 是常见做法ROCm 版本的 PyTorch 能跑但要注意torch.cuda.is_available()在 ROCm 下返回的也是 True别被名字骗了。如果是海光 GPU 或者麒麟系统 v10 这种国产化环境PyTorch 适配会更麻烦通常需要厂商提供定制 wheel不要直接 pip 装官网版本。先跑一段确认环境的代码别急着解压工程import torch print(torch version:, torch.__version__) print(cuda available:, torch.cuda.is_available()) print(cuda version:, torch.version.cuda) print(device count:, torch.cuda.device_count()) if torch.cuda.is_available(): print(device name:, torch.cuda.get_device_name(0))这段代码的作用是确认三件事PyTorch 装没装对、CUDA 能不能用、显卡认没认出来。torch.version.cuda如果是 None说明你装的是 CPU 版后面训练会慢到怀疑人生。device_count为 0 就是没识别到卡先解决驱动再谈模型。参数上没什么可调的但输出里的 CUDA 版本要和你nvidia-smi里的驱动版本对得上差太多会出现能 import 但一跑就报CUDA error: no kernel image is available。Python 和 PyTorch 的版本对应也要注意。常见组合是 Python 3.8 配 PyTorch 1.10 到 1.13Python 3.10 配 PyTorch 2.0 以上。如果你用 Anaconda 配置环境建议单独建一个 env别在 base 里折腾conda create -n ffdnet python3.9 -y conda activate ffdnet pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python numpy scikit-image tqdm这里cu118是 CUDA 11.8 对应的 wheel 源你要按自己驱动改。opencv-python和scikit-image是 FFDNet 工程里读图、算 PSNR 常用的tqdm看训练进度。装完之后再跑一次上面的确认代码确保新环境里 torch 能用。2.2 解压 ffdnet-pytorch.zip 后先看目录结构解压之后不要急着python train.py先tree或者ls -R看一眼。一个典型的 FFDNet PyTorch 工程大概长这样目录/文件作用你要不要动models/网络定义通常是ffdnet.py改网络结构才动data/数据集加载和噪声生成换数据格式要动train.py训练入口调参主要在这test.py推理入口跑自己的图改这里checkpoints/权重存放一般不动utils/图像读写、PSNR 计算基本不动如果解压出来没有requirements.txt就按 2.1 里那几个包手动装。如果models/ffdnet.py里 import 了torch.nn.functional as F但没用到不用管老工程常见冗余。先做一次 import 测试确认工程内部没有循环依赖或者路径写死的问题import sys sys.path.append(./) # 确保能 import 到工程根目录 from models.ffdnet import FFDNet net FFDNet(in_nc3, out_nc3, nc64, nb15) print(net)in_nc和out_nc是输入输出通道彩色图是 3灰度图是 1。nc是中间特征通道数原论文用 64。nb是卷积层数原版 15 层。这三个参数决定了模型大小改小能跑得快但降噪效果会掉。如果这一步报ModuleNotFoundError八成是路径问题检查你是不是在工程根目录下执行的。2.3 用一张图跑通推理确认权重能加载训练之前先跑推理这是最快的验证方式。假设checkpoints/里有一个ffdnet_gray.pth或者ffdnet_color.pthimport torch import cv2 import numpy as np from models.ffdnet import FFDNet device torch.device(cuda if torch.cuda.is_available() else cpu) net FFDNet(in_nc3, out_nc3, nc64, nb15).to(device) state torch.load(checkpoints/ffdnet_color.pth, map_locationdevice) net.load_state_dict(state) net.eval() img cv2.imread(test.jpg)[:, :, ::-1] / 255.0 # BGR 转 RGB 并归一化 img torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0).to(device) noise_sigma 25.0 / 255.0 noise_map torch.full((1, 1, img.shape[2], img.shape[3]), noise_sigma).to(device) with torch.no_grad(): out net(img, noise_map) out out.squeeze(0).cpu().numpy().transpose(1, 2, 0) out np.clip(out * 255.0, 0, 255).astype(np.uint8) cv2.imwrite(denoised.png, out[:, :, ::-1])关键点在noise_mapFFDNet 的噪声水平图是单通道的尺寸和输入图一样值就是 sigma/255。noise_sigma设成 25/255 表示你告诉模型「这张图大概有 sigma25 的噪声」。如果你不知道实际噪声水平可以多试几个值或者用工程里带的噪声估计脚本。map_locationdevice是为了避免权重在 CPU 上保存、GPU 上加载时报错。跑完看denoised.png如果比原图干净且没有明显涂抹感说明权重和网络结构对上了。3. 训练自己的 FFDNet数据、噪声图和 loss 怎么配3.1 数据集准备和噪声生成策略FFDNet 训练数据一般用 BSD400、DIV2K 或者自己攒的干净图。工程里data/目录通常有个dataset.py核心逻辑是读干净图、随机裁 patch、加高斯噪声、生成对应的 noise map。如果你要换自己的数据改__getitem__里的路径和裁剪尺寸就行。import random import numpy as np import torch from torch.utils.data import Dataset class DenoiseDataset(Dataset): def __init__(self, img_list, patch_size128, sigma_range(0, 50)): self.img_list img_list self.patch_size patch_size self.sigma_range sigma_range def __len__(self): return len(self.img_list) def __getitem__(self, idx): img cv2.imread(self.img_list[idx])[:, :, ::-1] / 255.0 h, w, _ img.shape if h self.patch_size or w self.patch_size: img cv2.resize(img, (self.patch_size, self.patch_size)) h, w, _ img.shape 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] sigma random.uniform(*self.sigma_range) / 255.0 noise np.random.randn(*patch.shape) * sigma noisy np.clip(patch noise, 0, 1) noisy torch.from_numpy(noisy.transpose(2, 0, 1)).float() clean torch.from_numpy(patch.transpose(2, 0, 1)).float() noise_map torch.full((1, self.patch_size, self.patch_size), sigma).float() return noisy, noise_map, cleansigma_range是训练时噪声水平的范围原论文用 0 到 50。如果你只关心某个固定噪声段比如 15 到 25就把范围收窄模型在这个区间会更准。patch_size128 是常见值显存不够就降到 64。noise_map 的尺寸是(1, H, W)不是(1, 1, H, W)因为后面网络里会把它和特征图在通道维拼接这个维度顺序错了会直接报 shape 不匹配。3.2 训练循环里 loss 和优化器的选择FFDNet 原论文用的是 MSE loss简单直接。优化器 Adam学习率 1e-3训练到收敛大概几十个 epoch。工程里train.py一般已经写好了但你要知道每个参数在哪import torch.nn as nn import torch.optim as optim criterion nn.MSELoss() optimizer optim.Adam(net.parameters(), lr1e-3) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) for epoch in range(50): net.train() for noisy, noise_map, clean in dataloader: noisy noisy.to(device) noise_map noise_map.to(device) clean clean.to(device) optimizer.zero_grad() out net(noisy, noise_map) loss criterion(out, clean) loss.backward() optimizer.step() scheduler.step() print(fepoch {epoch}, loss {loss.item():.6f})lr1e-3是起点如果 loss 震荡就降到 1e-4。StepLR每 10 个 epoch 乘 0.5这是比较稳的衰减策略。loss.item()只反映最后一个 batch想看整体趋势就在 epoch 内累加求平均。注意net.train()和net.eval()要切换虽然 FFDNet 没有 BN 和 Dropout但保持习惯没坏处。3.3 验证和 PSNR 计算别只看 loss训练时 loss 降不代表图好看要定期跑验证集算 PSNR。工程里utils/一般有calculate_psnrdef calculate_psnr(img1, img2): mse np.mean((img1 - img2) ** 2) if mse 0: return float(inf) return 20 * np.log10(1.0 / np.sqrt(mse))验证时把net.eval()打开torch.no_grad()包住噪声用固定 sigma 而不是随机这样 PSNR 才有可比性。我一般每 5 个 epoch 存一次权重命名带 epoch 和 PSNR比如ffdnet_epoch20_psnr32.5.pth后面挑最好的用。如果 PSNR 一直上不去先检查 noise_map 的 sigma 和实际加的噪声对不对得上这是最常见的翻车点。4. 避坑与排查ffdnet-pytorch 跑不起来时先看这几条4.1 报错 CUDA out of memory现象训练一开始就RuntimeError: CUDA out of memory或者跑几个 batch 之后爆显存。原因patch_size太大、nc或nb设太高、batch size 没降下来。FFDNet 虽然不大但 128 patch 加 64 通道 15 层显存占用也不低。解决先把 batch size 降到 4 或 2再把patch_size降到 64。如果还不行把nc从 64 降到 32nb从 15 降到 12。显存实在紧张就用torch.cuda.empty_cache()在 epoch 之间清一下但根本办法还是减模型或减输入。4.2 推理结果全黑或者全白现象跑完推理denoised.png要么全黑要么全白或者颜色完全不对。原因归一化和反归一化搞反了或者 BGR/RGB 通道顺序错了。OpenCV 读进来是 BGRPyTorch 里一般按 RGB 处理中间要转。另外out * 255.0之前要clip到 0 到 1不然溢出。解决按 2.3 里的代码检查[:, :, ::-1]有没有加np.clip有没有做。如果输出是灰度但输入是彩色检查in_nc和out_nc是不是设成了 1。4.3 加载权重报 key 不匹配现象load_state_dict报Missing key(s)或者Unexpected key(s)。原因网络结构定义和权重保存时的结构不一致常见于nc、nb改了但权重还是旧的或者用了DataParallel保存的权重带module.前缀。解决先打印state.keys()和net.state_dict().keys()对比。如果是module.前缀问题用state {k.replace(module., ): v for k, v in state.items()}去掉。如果是结构不一致要么改回原参数要么只加载匹配的部分。4.4 训练 loss 不降或者降了但图更糊现象loss 卡在某个值不动或者 loss 降了但验证集 PSNR 反而掉。原因学习率太大导致震荡或者噪声生成和 noise_map 不匹配模型学了个恒等映射。解决先把 lr 降到 1e-4 试。然后检查noise_map的值是不是和实际加的噪声 sigma 一致不一致的话模型收到的条件信息是错的。如果 loss 降但 PSNR 掉大概率是过拟合加数据增强或者早停。4.5 在 WSL 或国产化环境里跑不起来现象WSL 里torch.cuda.is_available()返回 False或者海光/麒麟环境下 import torch 就报错。原因WSL 需要装对驱动和 CUDA toolkitAMD 卡走 ROCm 还要额外配置。国产化环境通常需要厂商定制版 PyTorch。解决WSL 里先确认nvidia-smi能跑再确认 PyTorch 装的是 CUDA 版不是 CPU 版。AMD 卡确认 ROCm 版本和 PyTorch 版本匹配。国产化环境直接找厂商要适配好的 wheel别自己编译坑太深。5. 把 FFDNet 用起来从单图降噪到批量处理和 ONNX 导出跑通训练和推理之后真正要落地通常是两个方向批量处理自己的图或者把模型导出成 ONNX 方便部署。批量处理很简单把 2.3 的推理代码包一层循环遍历文件夹里的图每张图按固定 sigma 或者估计的 sigma 处理。我一般会写一个batch_denoise.py参数用 argparse 传这样不用每次改代码import argparse import os import cv2 import numpy as np import torch from models.ffdnet import FFDNet parser argparse.ArgumentParser() parser.add_argument(--input, typestr, requiredTrue) parser.add_argument(--output, typestr, requiredTrue) parser.add_argument(--sigma, typefloat, default25.0) parser.add_argument(--weights, typestr, defaultcheckpoints/ffdnet_color.pth) args parser.parse_args() device torch.device(cuda if torch.cuda.is_available() else cpu) net FFDNet(in_nc3, out_nc3, nc64, nb15).to(device) net.load_state_dict(torch.load(args.weights, map_locationdevice)) net.eval() os.makedirs(args.output, exist_okTrue) for name in os.listdir(args.input): if not name.lower().endswith((.png, .jpg, .jpeg, .bmp)): continue img cv2.imread(os.path.join(args.input, name))[:, :, ::-1] / 255.0 img_t torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0).to(device) sigma args.sigma / 255.0 noise_map torch.full((1, 1, img_t.shape[2], img_t.shape[3]), sigma).to(device) with torch.no_grad(): out net(img_t, noise_map) out out.squeeze(0).cpu().numpy().transpose(1, 2, 0) out np.clip(out * 255.0, 0, 255).astype(np.uint8) cv2.imwrite(os.path.join(args.output, name), out[:, :, ::-1])--sigma是你要告诉模型的噪声水平如果一批图噪声差异大可以先用简单的噪声估计比如拉普拉斯算子的方差算个大概再传进来。--weights指向你训练最好的那个权重。这个脚本跑起来就是python batch_denoise.py --input ./noisy --output ./clean --sigma 25。导出 ONNX 是为了脱离 PyTorch 环境部署比如在 C 或者推理引擎里跑dummy_input torch.randn(1, 3, 256, 256).to(device) dummy_sigma torch.full((1, 1, 256, 256), 25.0/255.0).to(device) torch.onnx.export( net, (dummy_input, dummy_sigma), ffdnet.onnx, input_names[input, sigma], output_names[output], dynamic_axes{input: {2: h, 3: w}, sigma: {2: h, 3: w}, output: {2: h, 3: w}}, opset_version11 )dynamic_axes让导出的模型支持任意尺寸输入不然只能跑 256x256。opset_version11兼容性比较好再高有些推理引擎不认。导出之后用onnxruntime跑一遍对比 PyTorch 输出误差在 1e-4 以内算正常。最后说个我自己的习惯每次改完网络结构或者训练参数先拿一张固定测试图跑一遍把 PSNR 和耗时记下来和上次对比。FFDNet 这种模型参数动一点效果就可能差不少没有基线对比很容易越调越乱。希望帮到你。本文还有配套的精品资源点击获取