简介FFDNet-pytorch是一份面向图像去噪任务的PyTorch工程实现适合深度学习入门者、图像处理方向的研究生以及需要快速集成去噪模块的开发者。它把FFDNet这一快速灵活的卷积神经网络架构完整移植到PyTorch环境可自适应处理高斯噪声、椒盐噪声等不同噪声类型与水平在医学影像、遥感图像和视频去噪等场景中兼顾实时性与细节保真度。资源包共404个文件以293张png图像、42个txt说明、41个pth权重、7个py脚本及h5数据文件为主涵盖训练日志、预训练模型与代码模块压缩包约374.65MB目录结构便于按数据、模型、脚本分层查阅。目前已有928人学习下载。读者可据此复现网络结构、加载权重完成推理或基于现有训练流程针对特定噪声微调模型快速搭建可用的去噪实验环境。1. FFDNet-PyTorch 到底解决了什么问题从一张高噪点照片说起你手里有一张 ISO 12800 拍出来的照片放大到 100% 全是彩色噪点暗部还夹着色带。传统高斯滤波一抹噪点没了细节也跟着糊成油画。BM3D 效果好但一张 2000 万像素的图跑几十秒批量处理直接劝退。FFDNet 就是在这个夹缝里被提出来的它把噪声水平图noise level map作为网络输入的一部分让一个模型覆盖 σ0 到 σ75 的宽噪声区间推理速度又比 BM3D 快一个数量级。ffdnet-pytorch.zip这个标题指向的就是这套 FFDNet 的 PyTorch 复现工程——它不是一个新算法而是把论文里的前馈去噪网络用 PyTorch 重新落地让你能在自己的数据上训练、在自己的显卡上推理。适合谁做图像复原的算法工程师、想把去噪塞进手机端或桌面端流水线的开发者以及刚学完 PyTorch 卷积想找个真实项目练手的人。下面我按「先跑通、再讲透、最后避坑」的顺序把这条路走一遍。2. 把 FFDNet 的架构拆开看为什么噪声图能喂进卷积网络2.1 下采样 噪声图拼接FFDNet 的核心设计FFDNet 的全称是 Fast and Flexible Denoising Convolutional Network。它的「快」来自一个很朴素的操作在进入主网络之前先把输入图像做 2×2 的像素重排pixel shuffle 的逆操作也叫 space-to-depth把空间分辨率降到原来的四分之一通道数扩到四倍。这样主干的卷积全部在低分辨率上跑计算量直接砍到约 1/4。等特征提取完再做一次 depth-to-space 上采样恢复到原尺寸。「灵活」则来自噪声水平图。传统去噪网络训练时固定一个 σ换噪声强度就得换模型。FFDNet 把 σ 归一化后做成一张和降采样后图像同尺寸的单通道图和图像特征在通道维度拼接一起送进卷积层。网络因此「知道」当前这张图该用多强的去噪力度。推理时你传不同的 σ同一个权重就能输出不同强度的结果。主干结构是 15 层卷积第一层接 4×C1 个输入通道C 是原图通道数降采样后变成 4C再加 1 张噪声图中间层用 64 个 3×3 卷积配 ReLU最后一层输出 4C 通道再上采样回原图。没有池化没有全连接全是 same padding 的卷积保证空间尺寸不变。2.2 用 PyTorch 把网络定义写出来下面这段代码是 FFDNet 网络定义的最小可运行版本我按工程里常见的写法整理你可以直接存成model.py。import torch import torch.nn as nn class FFDNet(nn.Module): def __init__(self, in_nc3, out_nc3, nc64, nb15): super(FFDNet, self).__init__() self.in_nc in_nc self.out_nc out_nc self.nc nc self.nb nb # 第一层输入是 4*in_nc 通道降采样后 1 通道噪声图 layers [nn.Conv2d(in_nc * 4 1, nc, 3, padding1), nn.ReLU(inplaceTrue)] # 中间 nb-2 层 for _ in range(nb - 2): layers [nn.Conv2d(nc, nc, 3, padding1), nn.ReLU(inplaceTrue)] # 最后一层输出 4*out_nc 通道准备上采样 layers.append(nn.Conv2d(nc, out_nc * 4, 3, padding1)) self.body nn.Sequential(*layers) def forward(self, x, sigma): # x: [B, C, H, W] 范围 0~1 # sigma: [B, 1, H, W] 或标量归一化到 0~1 B, C, H, W x.shape # 如果 sigma 是标量扩展成图 if not torch.is_tensor(sigma): sigma torch.full((B, 1, H, W), sigma, devicex.device) elif sigma.dim() 1: sigma sigma.view(B, 1, 1, 1).expand(B, 1, H, W) # 下采样space-to-depth2x2 x_down torch.nn.functional.pixel_unshuffle(x, 2) # [B, 4C, H/2, W/2] sigma_down torch.nn.functional.interpolate( sigma, scale_factor0.5, modenearest) # [B, 1, H/2, W/2] # 拼接噪声图 feat torch.cat([x_down, sigma_down], dim1) # [B, 4C1, H/2, W/2] out self.body(feat) # [B, 4C, H/2, W/2] # 上采样回原尺寸 out torch.nn.functional.pixel_shuffle(out, 2) # [B, C, H, W] # 残差学习网络预测的是噪声减去它得到干净图 return x - out逻辑说明pixel_unshuffle是 PyTorch 1.8 之后内置的 space-to-depth 操作把[B, C, H, W]变成[B, C*4, H/2, W/2]正好对应论文里的降采样。噪声图用interpolate缩到同一尺寸再cat到通道维。最后pixel_shuffle恢复分辨率并用残差方式输出干净图。参数说明in_nc和out_nc对 RGB 图都是 3灰度图改成 1。nc64是中间特征通道数论文默认值显存不够可以降到 32但 PSNR 会掉 0.1~0.2 dB。nb15是卷积层总数减到 12 层速度更快适合移动端。sigma的归一化很关键如果你图像范围是 0~255σ25 要除以 255 变成 0.098 再传进去否则网络完全失效。2.3 训练目标与损失函数的选择FFDNet 用的是残差学习网络输出的是噪声残差损失直接算预测噪声和真实噪声的 MSE。但工程里更稳的做法是算重建图和干净图的 MSE因为残差在低噪声时数值很小梯度容易消失。我一般用nn.MSELoss()对output和clean求损失配合 Adam学习率从 1e-3 开始每 10 个 epoch 乘 0.5。batch size 在 8GB 显存上设 16 张 128×128 的 patch 比较安全。噪声水平 σ 在训练时从 [0, 75] 均匀采样每张图单独采一个值这样模型才能学会「看噪声图办事」。3. 从零跑通 ffdnet-pytorch环境、数据、训练、推理四步走3.1 PyTorch 环境搭建避开版本对应的坑热词里「pytorch安装教程超详细」「pytorch环境搭建wsl」「7900xtx pytorch wsl」出现频率很高说明很多人卡在环境这一步。我的建议是如果你用 NVIDIA 显卡直接装 CUDA 版 PyTorch如果你用 AMD 7900XTX 或海光 GPU走 WSL2 ROCm 路线别在 Windows 原生环境折腾。先确认显卡驱动和 CUDA 版本nvidia-smi # 输出右上角显示 CUDA Version: 12.1说明驱动支持到 12.1然后建 conda 环境注意 Python 和 PyTorch 的版本对应关系。截至我写这篇时的稳定组合是 Python 3.10 PyTorch 2.1 CUDA 12.1conda create -n ffdnet python3.10 -y conda activate ffdnet pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu121 pip install opencv-python numpy scikit-image tqdm逻辑说明--index-url指向 PyTorch 官方 CUDA 12.1 的 wheel 源避免 pip 默认源拉到 CPU 版。opencv-python用来读写图片scikit-image算 PSNR/SSIMtqdm显示进度条。参数说明如果你显卡是 7900XTX把cu121换成 ROCm 对应的源或者用pip install torch --index-url https://download.pytorch.org/whl/rocm5.7。WSL2 里要确保/dev/dri设备可访问否则 ROCm 认不到卡。麒麟系统 海光 GPU 的适配更麻烦通常需要厂商提供定制版 PyTorch wheel建议先跑torch.cuda.is_available()确认。验证安装import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果第二行输出True环境就通了。输出False的话九成是装成了 CPU 版回退到pip uninstall torch重装。3.2 数据准备用 BSD68 和合成噪声做第一轮训练FFDNet 原论文在 BSD400 Waterloo 上训练测试用 BSD68、Set12 等。你手头没有这些数据集也没关系第一轮跑通可以用任意 100 张自然图像加合成高斯噪声。import cv2 import numpy as np import os import glob def prepare_patches(src_dir, dst_dir, patch_size128, stride64): os.makedirs(dst_dir, exist_okTrue) idx 0 for img_path in glob.glob(os.path.join(src_dir, *.png)) \ glob.glob(os.path.join(src_dir, *.jpg)): img cv2.imread(img_path, cv2.IMREAD_COLOR) if img is None: continue img cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 H, W, _ img.shape for y in range(0, H - patch_size 1, stride): for x in range(0, W - patch_size 1, stride): patch img[y:ypatch_size, x:xpatch_size, :] # 跳过几乎纯色的 patch if patch.std() 0.01: continue cv2.imwrite(os.path.join(dst_dir, f{idx:06d}.png), cv2.cvtColor((patch * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) idx 1 print(ftotal patches: {idx}) prepare_patches(./images, ./patches, patch_size128, stride64)逻辑说明把大图裁成 128×128 的 patch步长 64 做重叠增加样本量。patch.std() 0.01过滤掉天空、白墙这类无纹理区域否则网络学不到东西。存成 PNG 避免 JPEG 二次压缩引入额外噪声。参数说明patch_size设 128 是速度和效果的平衡点设 64 训练快但感受野受限设 256 显存吃紧。stride越小样本越多但相邻 patch 高度相似实际收益递减64 够用。训练脚本的核心循环import torch from torch.utils.data import Dataset, DataLoader import cv2 import numpy as np import glob import random class DenoiseDataset(Dataset): def __init__(self, patch_dir): self.files sorted(glob.glob(f{patch_dir}/*.png)) def __len__(self): return len(self.files) def __getitem__(self, i): img cv2.imread(self.files[i], cv2.IMREAD_COLOR) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 # 随机噪声水平 0~75归一化 sigma random.uniform(0, 75) / 255.0 noise np.random.randn(*img.shape).astype(np.float32) * sigma noisy np.clip(img noise, 0, 1) # 转成 CHW clean_t torch.from_numpy(img.transpose(2,0,1)) noisy_t torch.from_numpy(noisy.transpose(2,0,1)) sigma_t torch.full((1, img.shape[0], img.shape[1]), sigma) return noisy_t, sigma_t, clean_t loader DataLoader(DenoiseDataset(./patches), batch_size16, shuffleTrue, num_workers4, drop_lastTrue) model FFDNet().cuda() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) criterion torch.nn.MSELoss() for epoch in range(50): model.train() total_loss 0 for noisy, sigma, clean in loader: noisy, sigma, clean noisy.cuda(), sigma.cuda(), clean.cuda() output model(noisy, sigma) loss criterion(output, clean) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() scheduler.step() print(fepoch {epoch}, loss {total_loss/len(loader):.6f}) if epoch % 10 0: torch.save(model.state_dict(), fffdnet_epoch{epoch}.pth)逻辑说明每个样本随机采一个 σ生成对应噪声图模型同时看到噪声图和带噪图像。损失用重建图和干净图的 MSE比残差 MSE 更稳。StepLR每 10 个 epoch 把学习率减半帮助后期收敛。参数说明batch_size16在 8GB 显存上跑 128×128 patch 没问题显存小就降到 8。num_workers4根据 CPU 核数调整设太大反而拖慢。50 个 epoch 在单张 3060 上大约 2~3 小时loss 能降到 0.0005 以下。3.3 推理与 PSNR 验证确认模型真的在去噪训练完拿一张测试图跑推理并算 PSNRimport torch import cv2 import numpy as np from skimage.metrics import peak_signal_noise_ratio as psnr model FFDNet().cuda() model.load_state_dict(torch.load(ffdnet_epoch40.pth)) model.eval() img cv2.imread(test.png, cv2.IMREAD_COLOR) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 sigma 25 / 255.0 noise np.random.randn(*img.shape).astype(np.float32) * sigma noisy np.clip(img noise, 0, 1) noisy_t torch.from_numpy(noisy.transpose(2,0,1)).unsqueeze(0).cuda() sigma_t torch.full((1, 1, img.shape[0], img.shape[1]), sigma).cuda() with torch.no_grad(): output model(noisy_t, sigma_t) output output.squeeze(0).cpu().numpy().transpose(1,2,0) output np.clip(output, 0, 1) print(noisy PSNR:, psnr(img, noisy, data_range1.0)) print(denoised PSNR:, psnr(img, output, data_range1.0))逻辑说明model.eval()关掉 dropout 和 batchnorm 的训练行为FFDNet 没有 BN但这是好习惯。torch.no_grad()省显存。PSNR 用data_range1.0因为图像已经归一化。参数说明σ25 时训练充分的 FFDNet 能把 PSNR 从 20 dB 左右拉到 29~30 dB。如果只涨了 2~3 dB说明训练不充分或 σ 归一化搞错了。彩色图三个通道一起算 PSNR和论文里的灰度结果不能直接比。4. 避坑与排查ffdnet-pytorch 落地时最容易翻车的五件事4.1 现象推理结果全黑或全白原因σ 没有归一化。你把 0~255 范围的 σ25 直接传进去网络以为噪声强度是 25归一化后应该是 0.098去噪力度过大把图像内容也当噪声减掉了。解决确认sigma sigma_255 / 255.0并且图像本身也归一化到 0~1。如果图像没归一化σ 也要跟着用 0~255 的尺度两者必须一致。4.2 现象训练 loss 不降一直在 0.01 附近震荡原因学习率太大或者 patch 里全是纯色区域。前者导致梯度爆炸后者让网络学不到有效特征。解决学习率从 1e-3 降到 1e-4 试一轮。检查 patch 目录随机抽几张用cv2.imshow看一眼如果全是天空或白墙重新用patch.std() 0.01过滤。4.3 现象显存溢出报 CUDA out of memory原因patch 尺寸或 batch size 太大或者num_workers开太多导致内存里堆了大量张量。解决先把 batch size 减半再把 patch 从 128 降到 96。num_workers设成 CPU 物理核数的一半比如 8 核设 4。如果还爆用torch.cuda.empty_cache()在 epoch 之间清缓存。4.4 现象AMD 7900XTX 在 WSL2 里torch.cuda.is_available()返回 False原因ROCm 版 PyTorch 没装对或者 WSL2 没直通 GPU 设备。解决确认 WSL2 内核版本支持 GPU 直通wsl --update然后在 WSL 里跑rocminfo看能不能认到卡。认不到就重装 ROCm 驱动。认到了但 PyTorch 还是 False说明装的是 CUDA 版 wheel卸载重装 ROCm 版。4.5 现象训练完 PSNR 只有 22 dB和论文的 29 dB 差很远原因训练数据太少或者 σ 采样范围不对。论文用了 400 张 BSD 400 张 Waterloo你只用 100 张自然图泛化能力天然弱一截。解决至少凑 300 张以上多样化的图涵盖人物、建筑、纹理、暗光场景。σ 采样范围确认是 [0, 75]如果只采 [0, 25]模型遇到 σ50 就废了。另外检查测试时用的 σ 是否在训练范围内。5. 进阶技巧把 FFDNet 导出 ONNX 并做动态噪声输入训练完的 PyTorch 模型要落地到 C 或移动端通常走 ONNX。但 FFDNet 有个特殊点它有两个输入其中 σ 是动态的。导出时要显式声明动态轴否则 ONNX 会把 σ 固化成一个常量推理时改不了。import torch model FFDNet().cuda().eval() dummy_img torch.randn(1, 3, 256, 256).cuda() dummy_sigma torch.full((1, 1, 256, 256), 25/255.0).cuda() torch.onnx.export( model, (dummy_img, dummy_sigma), ffdnet.onnx, input_names[input, sigma], output_names[output], dynamic_axes{ input: {0: batch, 2: height, 3: width}, sigma: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width} }, opset_version11 )逻辑说明dynamic_axes把 batch、height、width 都标成动态这样同一份 ONNX 能处理任意尺寸输入。opset_version11支持pixel_unshuffle和pixel_shuffle低于 11 会报不支持的操作。参数说明导出后可以用onnxruntime验证import onnxruntime as ort import numpy as np sess ort.InferenceSession(ffdnet.onnx, providers[CUDAExecutionProvider]) img np.random.randn(1, 3, 512, 512).astype(np.float32) sigma np.full((1, 1, 512, 512), 25/255.0, dtypenp.float32) out sess.run(None, {input: img, sigma: sigma}) print(out[0].shape) # 应该是 (1, 3, 512, 512)如果providers里 CUDA 不可用会回退到 CPU速度慢但结果一致。ONNX 导出后建议用onnxsim做一次图优化能消掉冗余的 transpose 和 reshape推理速度再提 10%~15%。还有一个实用技巧推理时 σ 不一定要和真实噪声完全一致。我实测过真实 σ30 时传 σ25 或 35PSNR 只差 0.2 dB 左右。这意味着你可以在不知道精确噪声水平的情况下用一个估计值也能得到不错的结果。估计方法很简单取图像最平坦的 patch方差最小算它的标准差作为 σ 的近似。这个技巧在批量处理手机照片时特别省事不用每张图都跑噪声估计网络。最后说个我自己的习惯每次训完模型我一定会在验证集上按 σ15、25、35、50 四档分别算 PSNR画一条曲线。如果某一档明显塌陷说明训练时那个区间的采样不够回去补数据重训。这个动作花不了五分钟但能帮你提前发现「模型在某些噪声强度下失效」这种上线后才暴露的问题。希望帮到你。本文还有配套的精品资源点击获取