简介:面向深度学习与信号/图像去噪方向的研究者,这份压缩包提供了论文配套的TEM-NLnet深度去噪网络实现。项目围绕特定数据类型的噪声建模与去除展开,包含数据生成、预处理、噪声学习、模型构建与演示等关键模块,适合用于复现实验、理解网络架构或基于示例数据开展二次开发。包内共11个文件,以Python脚本为主(10个py),辅以说明文档(1个md),整体仅14KB,代码结构紧凑,便于快速定位核心逻辑与训练流程。目前已有135人学习下载,对于希望以轻量代码入门深度去噪或参考论文实现的读者具有一定参考价值。通过阅读源码与说明,可掌握从信号仿真、数据集构造到模型训练与去噪效果评估的完整链路,同时也能根据自身任务调整输入输出与网络参数,将其迁移至图像、时序等噪声处理场景。
1. TEM图像去噪为什么难做:这份NLnet论文资源能帮你省下哪些时间
透射电子显微镜(TEM)图像去噪一直是材料表征和生物结构分析里的老大难,不是因为算法少,而是因为TEM图像噪声不单纯——电子计数统计波动带来泊松成分,探测器读出又叠加上高斯成分,两者混合后,信号弱、细节尺寸小,对比度又低,常规去噪方法一上手就容易把晶格条纹或者单原子缺陷抹掉。TEM-NLnet是一个面向论文场景的深度去噪网络项目,核心思路是把非局部注意力机制引入卷积网络,让网络在去噪的同时保留TEM图像中的重复周期结构和局部细节。这个资源我拆完之后最大的感受是:它不只是给了一张网络图,而是把数据生成、训练配置、评估管线都串起来了,特别适合刚接触图像去噪、想快速跑通一套TEM噪声实验的同学,也适合需要和已有滤波方法做对比实验的科研党。前置要求不高,熟悉PyTorch基础数据流就行,有GPU最好,没有的话用小尺寸训练也能跑起来。
2. 噪声模型与网络设计:泊松-高斯混合噪声下NLnet为什么有效
2.1 TEM图像的噪声来源与混合噪声建模
TEM成像的几个噪声来源决定了你不能像处理普通照片那样简单套高斯去噪模型。首先是电子束的量子统计特性,每个像素位置上接收到的电子计数服从泊松分布,这也是TEM图像低剂量成像时噪声明显的根本原因;其次是探测器读出电路引入了高斯型热噪声,再加上荧光屏和相机本身的增益波动,最终观测图像其实是这两类噪声的叠加。常见做法是把退化过程写成:
y = α·Poisson(x/α) + η
其中 α 是增益因子,η 是均值为 0、方差为 σ² 的高斯噪声。这个式子看着简单,但直接决定了你训练数据的生成方式。很多复现失败的案例——包括我自己第一次跑的时候——都是直接把干净图像加上高斯噪声就完事,结果模型在模拟测试集上PSNR很好看,一放到真实TEM图像上就露馅,边缘过度平滑、背景颗粒感丢失。
import numpy as np def generate_tem_noise(clean, alpha=50.0, sigma=5.0): """ 泊松-高斯混合噪声模拟,输入输出均为float32数组,范围[0, 1] """ # 泊松分量:先放大再归一化,alpha越小噪声越强 poisson_part = np.random.poisson(clean * alpha) / alpha # 高斯分量:探测器读出噪声 gaussian_part = np.random.normal(0, sigma / 255.0, clean.shape) noisy = poisson_part + gaussian_part # 减去偏置,避免整体亮度偏移 noisy -= sigma / 255.0 return np.clip(noisy, 0.0, 1.0).astype(np.float32)这段代码里 alpha 是控制泊松强度的关键参数,alpha 越小,泊松噪声越明显,模拟的就是低剂量成像的场景;sigma 对应高斯分量,单位是按 8-bit 灰度范围换算的。注意最后把高斯噪声的均值减掉了,否则生成的带噪图像会整体偏亮,模型学到的残差里会混入一个固定偏置,导致推理时输出比干净图像暗一截。
2.2 NLnet网络结构:卷积主干加非局部注意力
NLnet的设计思路延续了深度去噪领域一个比较成熟的经验:网络主体用残差学习,预测的是噪声残差而不是直接输出干净图像。在这个结构上,作者加入了一个关键模块——非局部块(Non-Local Block)。和普通卷积只能看到局部感受野不同,非局部块会计算特征图上任意两个位置之间的相似度权重,然后用这个权重去聚合全图信息。放在TEM图像里,这个设计非常对症:晶体条纹、规则排列的原子柱往往在空间上呈现周期性,网络在一个位置学到干净的条纹模式,就能通过非局部机制去“帮忙”另一个噪声严重的同类位置,这是普通卷积做不来的。
我用PyTorch实现这个项目时,主干的非局部块代码大致是这样的:
import torch import torch.nn as nn import torch.nn.functional as F class NonLocalBlock(nn.Module): def __init__(self, in_channels, reduction=8): super().__init__() mid_channels = in_channels // reduction self.q_conv = nn.Conv2d(in_channels, mid_channels, 1) self.k_conv = nn.Conv2d(in_channels, mid_channels, 1) self.v_conv = nn.Conv2d(in_channels, in_channels, 1) self.softmax = nn.Softmax(dim=-1) self.out_conv = nn.Conv2d(in_channels, in_channels, 1) self.gamma = nn.Parameter(torch.zeros(1)) def forward(self, x): B, C, H, W = x.shape q = self.q_conv(x).view(B, -1, H * W).permute(0, 2, 1) # B, HW, C k = self.k_conv(x).view(B, -1, H * W) # B, C, HW v = self.v_conv(x).view(B, -1, H * W).permute(0, 2, 1) # B, HW, C attn = torch.bmm(q, k) # B, HW, HW attn = self.softmax(attn) # 归一化 out = torch.bmm(attn, v) # B, HW, C out = out.permute(0, 2, 1).view(B, C, H, W) return self.out_conv(out) * self.gamma + x理解这个模块要不要被绕进去,关键看 gamma 这个参数的设置。初始化为 0,意味着非局部块一开始等价于恒等映射,随着训练逐步打开,网络先稳住了基础卷积特征,再渐进地学习全局信息,这种做法能显著减少一开始训练时的不稳定。工程上还有一个细节:mid_channels 用 reduction 压缩了计算量,不然 H×W 级的 attention map 在 512×512 图像上是灾难级显存占用。实际部署时,我一般把非局部块放在编码器最深的一到两层,因为浅层特征空间信息保留得太强,非局部注意力算出的相似度大多还是近邻区域在起作用,反而浪费算力。
2.3 损失函数设计:L1和SSIM缺谁都有问题
这个项目损失函数定义上还是走主流路线,但组合方式有一些参考价值。纯MSE损失在频域上倾向于把图像推平,对边缘细节惩罚不足,尤其不适合结构敏感的TEM图像。项目中实际采用的是Charbonnier损失和MSSIM损失的加权组合。Charbonnier损失的定义方式是:
ρ(x) = sqrt(x² + ε²)
这是L1范数的可微逼近,对离群噪声点不敏感,又保持了优化过程中的梯度稳定性。MSSIM则是对不同尺寸窗口下的结构相似度做平均,抓住了人眼对图像结构的感知,但单独用又容易导致图像过度保持纹理、去噪不彻底。
def charbonnier_loss(pred, target): eps = 1e-3 return torch.mean(torch.sqrt((pred - target) ** 2 + eps ** 2)) def mssim_loss(pred, target): from torchmetrics.functional import multiscale_structural_similarity_index_measure return 1.0 - multiscale_structural_similarity_index_measure( pred, target, data_range=1.0 ) def combined_loss(pred, target): return charbonnier_loss(pred, target) + 0.5 * mssim_loss(pred, target)权重比例不是不能动,但0.5这个初始值我建议先别改。MSSIM的权重如果调到1.0以上,验证集上的PSNR会明显掉;调到0.1附近又会觉得图像细节欠保持。整个组合的表现对噪声水平有依赖性,高噪声场景下Charbonnier的收敛速度比MSE慢一点,但最终质量和平滑度都好不少。你复现的时候,可以照着这个思路在验证集上扫一遍权重,比盲目改网络结构涨点更快。
3. 数据准备与预处理:TEM图像切块、归一化与增强的落地做法
3.1 训练数据集:合成噪声对与真实验证集缺一不可
做TEM去噪的深度模型,最大的软肋就是干净图像不好找。真实TEM图像里很难获得严格无噪声的GT,所以主流的落地路线是在干净的高剂量图像或标注过的合成图像上人为加噪声,构造训练对。某个实验室复现这个项目时,所用的干净图像是从某公开显微图像库筛选出来的高信噪比TEM图,用2.1节那个generate_tem_noise函数按不同alpha、sigma组合生成带噪图像。除合成训练对之外,一定要留一小部分真实低剂量TEM图像作为人工验证集,并且带着它们参与下采样可视化,否则模型在真实场景是否可用完全没有参照。
数据划分上,我的习惯是不按图像随机分。TEM图像同一视场的相邻区域结构高度相似,如果划分不当,训练集和测试集出现重叠视野,会虚高PSNR。最好的做法是按照样本或采集批次划分,确保同一张原始大图的所有切块只落在同一个集合里。这个文件组织方式虽然看起来只是目录设计,但对论文实验的严谨性影响很大。
data/ ├── train/raw/ # 干净TEM原图,不切块 ├── train/noisy/ # 模拟噪声对照,文件名与raw一一对应 ├── val/real_low_dose/ # 真实低剂量图像,无GT ├── val/synthetic/ # 合成带噪图+干净图对 └── test/ # 最终报告PSNR的固定测试图3.2 切块与归一化:别把整张图的对比度差异压掉
TEM图像的灰度响应和常见摄影图像差别很大,像素值通常以16位深度存储,直接读进来缩放到0-1再训练容易翻车。原因是整张尺寸大的TEM图往往存在亮度梯度,不同位置的局域对比度变化明显,对整图做全局归一化之后,暗区域的噪声方差会被压缩得很小,模型本质上没有见过噪声显著但细节微弱的输入。
我一般会先做局部的滑动窗口归一化,或者在切块阶段对每个patch独立做per-patch的标准化。要注意per-patch标准化会抹掉块之间的绝对灰度差异,所以需要保证训练时输入网络的信息同时包含局域相对变化和全局全局亮度。常见做法是把图像先按全局统计量线性映射到0-1,再对每个patch计算局部均值方差,拼接成双通道输入送进网络:
def pack_patch(patch_01): """ patch_01: float32, [0,1],单张patch 返回双通道输入:原值通道 + 局域标准化通道 """ local_mean = patch_01.mean() local_std = patch_01.std() + 1e-6 local_norm = (patch_01 - local_mean) / local_std return np.stack([patch_01, local_norm], axis=0).astype(np.float32)这个trick不是网络结构的一部分,但很多论文实验的稳定涨点都来自这类数据层面的设计。局域标准化通道相当于一个适应性子层,帮助卷积核更快地响应不同亮度区域的噪声结构。切块尺寸一般用128×128或256×256,要保证网络最深层特征图至少还剩8×8,不然全局信息不够。128×128适合初期调参,256×256适合最终实验,两者都可以在推理阶段用重叠滑窗拼回原图。
3.3 数据增强:旋转翻转之外,噪声水平也要跟着变
TEM图像不像自然图像有方向和语义约束,随机90度旋转、水平翻转、垂直翻转都是安全的增强方式,不会产生“翻转后含义改变”的问题。但要特别注意,增强的顺序必须放在加噪声之前:如果先翻转再加噪声,等价于对同一块内容的不同噪声实现做增强,扩大了有效样本量;反过来则只是重复使用相同噪声模式,提升很有限。
在噪声层面,也可以在训练过程中动态生成每个batch的噪声水平。比如每个batch的alpha和sigma从某个预设范围内随机采样,这样模型见到的是一个噪声强度区间而不是单一噪声档位,实际推理时握着一点泛化余量。代码实现其实很省事,generate_tem_noise的一个参数从固定值改成随机采样即可,整体训练时间几乎没有增加,模型却不再迷恋某个特定噪声强度下的统计假设。
4. 训练配置与调参:从优化器到学习率策略的完整参考
4.1 优化器与初始超参数
这份资源里的实验配置,我拆解出来对比了一下自己的复现记录,发现它采用的是一套很实用的基线设置,适合作为起点而不是终点。优化器主流选择是Adam,默认beta值在去噪任务里通常够用,初始学习率取1e-4是一个兼顾稳定性和收敛速度的档位。SGD在TEM去噪这种回归任务上也能跑,但收敛曲线上抖动明显,调到合适的momentum和weight decay需要额外花时间,起步阶段不推荐。
裁块尺寸和batch size是直接挂钩的。256×256的输入,在单张24G显存的设备上,batch size上限通常在8到12之间;128×128的输入,batch size放宽到16问题不大。如果显存卡在8G到12G这个区间,优先降batch size到4到6,配合梯度累积模拟较大batch,而不是去裁小输入尺寸,否则感受野缩水带来的细节损失比batch size带来的统计噪声更严重。
一个值得记录的基础参数表:
| 参数 | 推荐值 | 备注 |
|---|---|---|
| 输入patch | 128×128 / 256×256 | 调参用128,实验用256 |
| batch size | 8(256px)/ 16(128px) | 小显存用梯度累积 |
| 初始学习率 | 1e-4 | 2e-4偏激进,容易让非局部块的比例崩掉 |
| 优化器 | Adam (betas=0.9, 0.999) | 高阶矩偶尔手动clip到1e-8 |
| 总轮数 | 100~200 | 50轮后再评估SSIM才有区分度 |
| 权重复选框 | 验证集PSNR+SSIM联合排序 | 只看PSNR会选到过度锐化的权重 |
4.2 训练循环代码框架
训练脚本的骨架并不复杂,但有几个细节值得写进代码里。首先,模型输出的残差加上输入才是预测的干净结果,所以loss要作用在重建后的完整图像上,GC作用于输出图本身;其次是在每个epoch结束时,用验证集噪声图做一次完整的前向推理,计算固定batch上的PSNR和SSIM,不能只看训练loss下降就认为模型在学。
import torch from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100) for epoch in range(start_epoch, total_epochs): model.train() for batch in train_loader: noisy, clean = batch["noisy"].cuda(), batch["clean"].cuda() optimizer.zero_grad() with autocast(): residue = model(noisy) pred = noisy + residue loss = combined_loss(pred, clean) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() scheduler.step()这段代码里最需要理解的是clip_grad_norm_这一行。不要省略。非局部模块的注意力权重在训练初期容易出现较大的梯度尖峰,如果不加梯度裁剪,训练到第10到20个epoch的时候loss可能突然跳到NaN。max_norm取1.0是一个比较稳妥的折中值,既能抑制尖峰,又不会明显拖慢收敛。顺带说明,mixed precision不是必须项,但配合gradscaler能省下近一半的显存,把256×256的batch size从4提升到8,这在实际实验中往往就是拟合质量的分水岭。
4.3 学习率调整与收敛判据
项目给出的默认策略是CosineAnnealing,最大训练轮数T_max设为100,这是目前去噪任务里性价比最高的做法。StepLR按epoch衰减(比如每30轮乘以0.5)也能用,但最后结果通常比cosine schedule低0.1到0.3个dB的PSNR。值得提醒的是,Pytorch里CosineAnnealing的T_max应该与总轮数对齐,不要设成和训练轮数无关的固定值,否则学习率退火节奏错位,模型在最优学习率区间停留时间不够。
判断模型是否收敛,我个人的经验是:训练初期前10个epoch的PSNR增长速度最快,20到50轮之间速度放缓,50轮之后基本进入缓慢爬坡阶段。如果验证集PSNR在20轮内完全不动,问题大概率出在数据生成上,而不是模型容量不足;如果训练loss持续下降但验证集指标回退,说明过拟合已在发生,此时应该回头检查数据增强范围是不是太窄,而不是增加正则项——因为去噪任务里过拟合的主要来源是噪声模式重复,而不是图像内容重复。
5. 实践避坑:四个导致模型翻车的常见问题与排查
5.1 模拟噪声参数与真实TEM噪声不匹配
现象:模型在合成测试集上PSNR超过30dB,迁移到真实低剂量TEM图像后图像变得光滑,原本清晰的碳支持膜背景纹理被抹成一片均匀的灰色,原子柱边缘也模糊了。
原因:训练时alpha和sigma只在单一档位上采样,而真实TEM噪声随着加速电压、束流强度和相机曝光时间变化很大,单一噪声参数覆盖不了实际分布的方差范围。
解决:把alpha和sigma改成训练时随机采样,比如alpha从20到80,sigma从0到10之间采样,这样模型看到的是一个噪声分布族而不是一个点。进阶方案是微调阶段混入少量真实带噪图像,用一个不带GT的一致性损失做约束,让模型输出的结构特征在相邻帧间保持一致,这种半监督微调对真实场景泛化的提升非常直接。
5.2 训练中途loss变NaN,非局部块是重灾区域
现象:训练进行到第15个epoch时,loss从0.002突然跳变成nan,重新从checkpoint恢复后同样在相近位置崩掉。
原因:非局部块中的注意力logits没有做数值稳定性处理。当特征图某些通道的响应值过大时,bmm计算出的相似度矩阵会出现巨大数值,经过softmax后梯度爆炸。这是这类结构任务里最容易翻车的地方,自注意力族模块的通病。
解决:除了在训练代码里加梯度裁剪,非局部块内部的attention计算也要改进。我习惯在torch.bmm得到相似度矩阵之后,对它做减最大值(sub max)操作再进行softmax,原理是利用softmax的平移不变性压缩数值范围。这个改动对最终指标几乎无影响,但能彻底规避NaN复现的随机性。
5.3 BatchNorm在小batch下抖动剧烈
现象:在两张24G卡上分别跑同一个配置文件,一张卡训练出的模型PSNR比另一张高出0.4dB,换回同一张卡复现时结果又不一致,看起来像是没有固定随机种子。
原因:BatchNorm的统计量在小batch(小于8)下估计方差大,训练和推理时使用的统计值不一致,导致模型性能出现明显波动。TEM去噪的patch尺寸通常较大,batch size往往卡在4到8之间,这就正好踩中BatchNorm的雷区。
解决:把主干里所有的BatchNorm替换为InstanceNorm或GroupNorm。TEM图像的噪声虽然有空间相关性,但我们并不需要依赖batch统计数据去刻画噪声分布,InstanceNorm每个样本独立归一化,训练推理行为一致,还省掉了对batch size的敏感性。我在实际替换后,两张卡上的结果差异缩小到0.05dB以内,训练也更稳定。
5.4 滑窗推理输出出现网格状拼接痕迹
现象:把大尺寸TEM图切成256×256重叠patch逐个推理,再拼回原图时,patch与patch衔接处出现明显的亮度断层,形成网格状伪影。
原因:拼图时只是简单把重叠区域直接取平均,或者每个patch独立做了归一化,导致相邻patch重叠区像素值不一致。TEM图像的噪声和细节在空间上变化快,直接平均会产生接缝。
解决:拼图时引入权重窗口,重叠区域按像素位置乘以线性/汉明权重再累加,同时确保所有patch使用全局一致的归一化参数,而不是每个patch单独估计。代码层面的要点是维护一个累计权重图,最终输出为加权累加除以累计权重。这种方式可以完全消除可见的拼接网格,对于要求高保真的论文配图来说是必备细节。
6. 评估与可视化:不只看PSNR,怎么判断去噪是否保住结构
拿到一个能稳定输出的模型之后,评估层面的工作对论文和工程结论的帮助同样重要。PSNR作为指标有明确短板——它对全局像素差值友好,但完全不在意空间结构。所以我的评估流程一定会做三件事:单独算SSIM和FSIM指标,生成残差图切片,以及做频谱对比验证高频细节有没有被过量抑制。只有在三个维度上都合理的模型,才是真正适合TEM图的去噪模型。
残差图的观察重点是去噪结果减干净图的差值图,它应该表现为无结构的随机噪声。如果残差图里隐约能看到晶体条纹或原子排布的轮廓,说明模型把信号从输入中“抹掉”了一部分——这在PSNR上可能只是轻微下降,但结构信息损失是肉眼可见的。频谱对比方面,傅里叶变换后的高频环或点阵谱,应当在去噪后保持清晰的峰值,而不能变成磨砂状的一片;如果高频能量整体塌陷,那无论PSNR是30还是31dB,这套模型都是过平滑的。
实践中,我通常会在一个测试集上同时输出这三个评估结果,制成一张横向对比图。这个习惯对我自己帮助很大:曾经有一版epoch更后、PSNR更高的权重,正是在残差图中暴露出明显条纹轮廓,我才发现它其实已经陷入了轻微的过拟合状态,换回早几十轮次checkpoint后,真实TEM图上的细节表现反而更好。有效验证了仅仅靠验证集指标选权重的机制是不可靠的,尤其是当任务涉及高频纹理和结构保真时。
从那以后,我每次训练完都会强制走一遍这个评估管线——算指标、看残差、看频谱,三者一个不少。因为跑过太多看起来PSNR不错、实际不能使用的模型,能避免这些翻车点,才是这套管线最大的价值。希望这个项目模板和调试记录也能帮你少走一些弯路。
本文还有配套的精品资源,点击获取