TEM图像去噪实战：NLnet与Non-Local机制详解

简介这份资源是围绕论文《TEM-NLnet: A Deep Denoising Network for T》整理的演示项目面向深度学习研究者、论文复现者以及关注信号/图像去噪的开发者。项目聚焦TEM-NLnet深度去噪网络用Python完整实现了数据生成、预处理、模型构建与训练推理流程有助于理解去噪自编码器、卷积或递归结构在实际噪声去除中的工作原理。压缩包共11个文件以10个Python脚本为主覆盖数据生成、工具函数、模型定义含GN模块以及去噪/噪声学习训练入口另有1份README说明文档整体体积仅14KB轻量紧凑方便快速阅读、修改和调试。目前已有135人学习适合作为论文配套代码的入门参考。通过这份资料读者可以梳理TEM-NLnet的数据流组织方式直接复用其数据生成与预处理脚本并依据自身数据替换输入、调整网络结构迁移至图像去噪、时间序列清理等任务为后续实验对比提供可运行的基线实现。1. TEM 图像去噪为什么难NLnet 是什么、能帮你解决什么拿到一张低剂量透射电子显微镜图像信噪比低到连晶格条纹都看不清时常规的滤波和普通 CNN 去噪很容易把细节抹成一片塑料质感。目录里的 TEM-NLnet 正是一个论文配套的深度学习去噪项目核心思路是用 Non-Local 机制捕捉电镜图像里大量重复出现的结构模式在压制噪声的同时保住纤维、晶格、颗粒边界这些真正有用的高频信息。它能解决的核心问题是电子束敏感的样品只能低剂量成像而低剂量图像噪声大有了这套网络和配套训练流程你可以在合成噪声数据上把模型练出来再迁移到真实电镜图像上做降噪。适合正在做电镜图像处理、论文复现或者手里攒了一批低信噪比 TEM 数据却不知道怎么下手的人。这篇笔记直接把解压、环境、数据、训练、排坑这一条链路讲完你照着做就能跑通。2. 解压后第一步看目录、搭环境、跑通一次最小推理2.1 压缩包解开后先别急着运行先看这三样东西论文配套代码和工业项目差距很大作者往往只保证在固定环境里能跑不会替你考虑依赖冲突。拿到后我会依次打开 README、requirements.txt 和主配置。README 里通常写了数据格式、预训练权重是否存在、训练入口命令requirements.txt 决定你是用 pip 直接装还是需要手动处理 PyTorch 版本配置则决定了一切数据路径在哪里改。这类去噪项目最常见的目录结构大概是这样供你对照着找入口TEM-NLnet/ ├── README.md ├── requirements.txt ├── configs/ │ ├── train.yaml │ └── test.yaml ├── data/ │ ├── train/ │ └── val/ ├── models/ │ ├── nlnet.py │ ├── blocks.py │ └── loss.py ├── utils/ ├── train.py ├── test.py └── infer.py不太可能每个文件都齐全但 train.py 和 test.py 基本是标配。如果你的压缩包里有预训练权重一般是 .pth 或 .pt 结尾那就先走一遍 test 或 infer如果只有模型代码没有权重就得从合成数据开始自己训练。这里有一个经验不要一上来就打开 model 文件读网络结构先去跑通 README 里的最小示例跑通了再回来看结构思路会顺很多。2.2 建独立环境PyTorch 版本和 CUDA 匹配是第一道坎TEM 去噪模型基本都是 PyTorch 写的环境搭错了后面全是坑。我会为这个项目单独建一个虚拟环境不往基础环境里乱装东西。常见做法是先看 requirements.txt 里锁没锁 PyTorch 版本如果没锁就按自己显卡驱动能支持的 CUDA 版本选 PyTorch而不是越新越好。cd TEM-NLnet python -m venv venv source venv/bin/activate # Windows 上用 venv\Scripts\activate pip install -r requirements.txt装完先验证一下 PyTorch 能不能正确调用 CUDA这一步省略的话常常会训练到一半才发现模型跑在 CPU 上速度慢几十倍也不报错。python -c import torch; print(torch.__version__, torch.cuda.is_available())如果你的 test.py 依赖特定版本的数据读取库比如要求 tifffile而 requirements.txt 里没有那就手动补装。这里建议优先用 tifffile 或 imageio不要用 OpenCV 读电镜图原因见后面第 5 章的坑简单说就是 OpenCV 默认按 uint8 读16 位 TIFF 的动态范围直接断送。2.3 跑通最小推理一张图验证整条链路模型代码写的再好跑不通就是零。我一般会先找一张真实的低剂量 TEM 图只要一张放进一个 samples 目录然后运行推理脚本。如果你有预训练权重python test.py \ --cfg configs/test.yaml \ --input samples/tem_low_dose.tif \ --weight checkpoints/tem_nlnet_best.pth \ --output outputs/tem_denoised.tif如果压缩包没有 test.py 而是统一用 train.py 带 phase 参数就去看 README 里的对应写法和实际入口。跑通之后把输出的 tif 和输入放在一起对比一下噪声是否明显下降晶格条纹有没有被磨平。这一步的意义是确认网络、数据流、权重加载全链路正常之后再进入训练环节。如果这一步都卡住优先看报错里涉及键名的 KeyError多半是权重训练时用的配置键和 test.yaml 不一致需要把两者对齐。3. 造训练数据TEM 噪声建模与 patch 采样3.1 TEM 噪声为什么不能直接套高斯普通照片去噪常常假设噪声是高斯分布这个假设在 TEM 上会翻车。透射电镜图像的底层噪声是电子计数的泊松过程尤其在低剂量条件下每个像素的计数可能只有几十到几百泊松分布会明显偏离高斯噪声方差和信号强度成正比。换句话说信号越强的地方噪声越大这是和高斯噪声最大的区别。再加上读出噪声和探测器响应实际图像噪声是信号相关的泊松加上一个相对固定的高斯尾。如果拿高斯假设去合成训练数据模型学出来的分布和真实 TEM 噪声对不上合成指标再漂亮一上真实图就露馅。所以做训练数据时我一般会按泊松-高斯混合模型来合成噪声对把干净图当成电子计数归一化乘以一个剂量系数后做泊松采样再叠一个很小的读出噪声。用这个方法你不需要真正去低剂量拍一整套数据也能让模型先学会基本的噪声压制。import numpy as np def add_tem_noise(img, dose300.0, read_noise_sigma2.0, clipTrue): 模拟低剂量TEM图像噪声。 img: 0~1范围的干净图单通道float64 dose: 平均电子计数越小噪声越大常用范围20~800 read_noise_sigma: 读出噪声标准差单位是灰度级/255 img np.asarray(img, dtypenp.float64) img_norm np.clip(img, 0.0, 1.0) photon_count img_norm * dose noisy_count np.random.poisson(photon_count).astype(np.float64) noisy_img noisy_count / dose noisy_img noisy_img np.random.normal(0, read_noise_sigma / 255.0, sizeimg.shape) if clip: noisy_img np.clip(noisy_img, 0.0, 1.0) return noisy_img关键参数就是 dose 和 read_noise_sigma。dose 控制泊松噪声强度等于模拟不同电子剂量下的成像质量read_noise_sigma 控制探测器底噪。这两个参数决定了训练数据的噪声分布范围直接影响模型在真实低剂量数据上的表现。后面第 4 章会讨论剂量范围怎么配这里先知道生成一对干净-带噪样本只需要这几行就够了。3.2 patch 采样决定训练稳定性和显存消耗训练去噪网络时很少把整张大图直接喂进去常见做法是切成 patch。TEM 图一般 2048x2048 甚至更大整图进网络显存立刻爆炸而且一张图的梯度更新太粗。先把数据切成 patch再按批训练这是更稳的方案。import numpy as np import tifffile def extract_patches(img, patch_size128, stride64): h, w img.shape[:2] patches [] for y in range(0, h - patch_size 1, stride): for x in range(0, w - patch_size 1, stride): patches.append(img[y:ypatch_size, x:xpatch_size]) return np.stack(patches)patch_size 和 stride 是两个直接决定训练效果的参数。patch_size 推荐从 128 开始太小比如 64会让 Non-Local 模块找不到足够的空间关系去噪容易发糊太大则显存开销猛增。stride 推荐取 patch_size 的一半比如 128 的 patch 用 64 的步幅保证相邻样本之间有重叠但不冗余。如果你有几张干净大图切完通常能得到几千个 patch这个量级对去噪网络训练已经够用。建议在切片之后做一次归一化用 2% 和 98% 分位数做线性拉伸而不是简单 min-max 拉伸因为电镜图像里经常有极亮的污染点min-max 会把正常动态范围压得很扁。3.3 伪 GT 的真实陷阱别拿粗暴去噪结果当干净标签做训练数据还有一个绕不开的难题真实 TEM 图像很难有真正的干净参考图。常见做法有两种一是拍同一视野的高剂量图当作干净标签低剂量图当输入配对训练二是用 BM3D 这类传统去噪把低剂量图洗一遍当作伪 GT 来训练。第一种效果好但耗时第二种看起来省事隐患很大。如果直接拿 BM3D 的结果当标签模型学到的其实是「BM3D 的特征 真实噪声残留」那些伪影会被当成正常结构学习训练出的网络在真实图上会输出一种「非自然」的平滑感。更隐蔽的是伪 GT 本身已经模糊了晶格细节模型学到的映射会把细节也一并抹掉。我的习惯是初始阶段可以用少量伪 GT 作为 pretraining之后再搭配真实低-高剂量配对数据做 fine-tune。如果条件允许直接跳过伪 GT工作量多但结果干净得多。4. NLnet 核心设计与三个必调参数4.1 Non-Local 机制为什么适合 TEM 图像普通卷积的感受野是局部的一层的有效范围可能只有几十个像素虽然堆叠层数能扩大感受野但代价是信息逐层衰减细微的高频结构在传递中丢失。Non-Local 的核心思想是让每个位置的输出由全图范围内相似位置的特征共同决定用所有和它相似的信息来辅助去噪。TEM 图像中大量存在周期性纹理、重复的颗粒、平行的纤维束这些结构在图上相距很远但内容高度相似恰好是 Non-Local 最擅长利用的。具体做法是对每个像素位置计算一个注意力权重权重由查询位置和候选区域的特征相似度决定相似度高的位置赋予更高权重。用公式来说输出是注意力权重对全图特征做加权平均的结果等价于给网络装了一个「全局查找相似结构」的能力。相比普通卷积这种机制的代价是计算量显著增加所以实际实现里一般会在低分辨率特征图上做 Non-Local或者把注意力计算限制在一个区域内。4.2 一个可落地的 Non-Local 模块实现如果压缩包里的 nlnet.py 读起来费劲或者你想自己验证这个机制可以先跑一个简化的 NonLocalBlock确认它对噪声图的效果。下面的实现采用了常见的降维策略把通道数从输入的大通道数压到 mid_ch大幅降低注意力矩阵的内存占用。import torch import torch.nn as nn class NonLocalBlock(nn.Module): def __init__(self, in_ch, mid_ch8): super().__init__() # mid_ch 是注意力空间的维度设得越小越省内存但太小时表达力受限 self.q nn.Conv2d(in_ch, mid_ch, kernel_size1) self.k nn.Conv2d(in_ch, mid_ch, kernel_size1) self.v nn.Conv2d(in_ch, in_ch, kernel_size1) self.proj nn.Conv2d(in_ch, in_ch, kernel_size1) def forward(self, x): B, C, H, W x.shape Q self.q(x).view(B, mid_ch, H * W).permute(0, 2, 1) # B, HW, mid_ch K self.k(x).view(B, mid_ch, H * W) # B, mid_ch, HW V self.v(x).view(B, C, H * W) # B, C, HW attn torch.softmax(Q K / (K.size(1) ** 0.5), dim-1) # B, HW, HW out (attn V.permute(0, 2, 1)).permute(0, 2, 1).view(B, C, H, W) return self.proj(out)这段代码里 Q 和 K 的矩阵乘法会生成 HW 乘 HW 的注意力图分辨率一大内存会指数级上升所以 mid_ch 和特征图分辨率都需要控制。除以 K.size(1) 的开方是为了让 softmax 前的数值不会因为维度增大而过大这是一个稳定的缩放因子。实际项目中 NLnet 会在编码器-解码器的中间层插入这种模块而不是在网络最前端直接做因为浅层特征图分辨率太高计算全图注意力不现实。4.3 训练前要定死的三个关键参数训练一个 TEM 去噪网络最重要的参数不是网络结构本身而是下面这三个它们决定了模型能不能稳定收敛、能不能在真实数据上泛化。第一个是 patch_size我建议至少 128这条前面已经说过它直接影响 Non-Local 能看到多大的空间上下文。第二个是训练噪声的 dose 范围这一步最容易被忽略如果你只在 dose800 下合成数据训练模型遇到 dose50 的真实低剂量图像时噪声水平超出训练分布表现会明显退化。我的习惯是让 dose 在一个宽范围内随机采样比如 20 到 800 之间取对数均匀分布让模型见过从极低剂量到接近正常成像的各种情况。第三个是 loss很多论文用 L2但 L2 在 TEM 去噪中容易让输出过度平滑边缘发虚。相比之下 L1 或 Charbonnier 损失对边缘保持更友好是更稳的选择。5. 训练评估与常见坑5 条排查路径5.1 训练阶段的三个高频坑不收敛、OOM、纹理消失第一个坑是 loss 下降很慢甚至震荡。现象是训练曲线忽上忽下偶尔还出现 NaN。原因通常是学习率过大或者输入图像没有归一化到 0 到 1 的范围导致梯度幅度差异巨大。解决方法是先把输入-输出的数据分布统一到 0~1再把学习率从默认的 3e-4 调低到 1e-4并加上前 5 个 epoch 的 warmup。如果用的是 AdamW一般不需要额外的 weight decay 调优0.01 默认即可。第二个坑是显存不足这在 Non-Local 网络里尤其常见。现象是 batch_size 改成 1 依然 OOM。原因是 Non-Local 在较大特征图上生成了 HW 乘 HW 的注意力矩阵。解决方法是把 Non-Local 块放到分辨率已经缩减一半的编码器中间层同时把 patch_size 从 128 调整到 96 或 64。如果还不行可以用梯度累积把梯度攒几个 step 再更新效果等效于增大 batch_size 但显存不变。要注意的是 patch_size 缩小后要重新验证去噪质量不要为了省显存牺牲过多空间上下文。第三个坑是训练完去噪效果好但纹理细节被磨平了。现象是晶格条纹变成模糊的亮暗交替颗粒边界不清晰。原因很可能是 loss 用了纯 L2或者下采样次数太多导致高频信息在瓶颈层被压缩。解决方法是换 Charbonnier 损失同时检查网络的 stride 设置下采样不要超过三次。贴一个可以直接用的 Charbonnier lossimport torch import torch.nn as nn class CharbonnierLoss(nn.Module): def __init__(self, eps1e-3): super().__init__() self.eps eps def forward(self, pred, target): diff pred - target return torch.mean(torch.sqrt(diff * diff self.eps ** 2))eps 参数控制损失函数在零附近的平滑程度取 1e-3 时既保持了接近 L1 的鲁棒性又不会像 L1 那样在零点不可导。梯度流通更稳定收敛也更平滑。5.2 评估阶段的两个高频坑数据位深和真实泛化第四个坑很隐蔽经常在最后评估时才爆炸模型在合成数据上 PSNR 很高放到真实 TEM 图上却效果很差。原因是合成数据里的噪声分布和真实数据不一致最常见的是模拟噪声时没有考虑探测器响应或者剂量范围设定得太窄。解决方法是先检查真实低剂量图的噪声标准差统计一块均匀区域的灰度标准差再用这个实测值去调合成分散参数。如果不想做精细统计就用一个粗暴但有效的办法把 dose 范围从 10 到 2000 全部放进训练集让模型能覆盖超低到正常剂量的完整区间再在真实图上做少量 fine-tune。第五个坑是 16 位 TIFF 被读成 8 位整个图动态范围塌缩。现象是你的网络输入 0~1 的范围但读入的数据实际上是被截断过的 0~255 灰度噪声被当作信号去噪效果完全不可用。原因是 OpenCV 的 imread 默认按 8 位读取16 位图的 12 位或 14 位有效信息全部丢失。解决方法是统一用 tifffile 读取并且按实际位深做归一化。import tifffile import numpy as np img tifffile.imread(samples/tem_low_dose.tif) img img.astype(np.float64) # 假设这张图是12位有效数据最大值约4095 img_norm img / 4095.0 img_norm np.clip(img_norm, 0, 1)位深信息可以从 TIFF 文件的 metadata 里读也可以在预处理阶段先打印图像数组的最大值反推。这一步不做好之前的所有训练工作都会白费数据范围错了整个去噪就是另一个问题。另外评估 PSNR 时也要注意如果一张图是 12 位、一张是 16 位转出来的千万不能直接把它们放进同一个归一化尺度否则指标会被明显拉低。5.3 快速定位问题的排查顺序遇到训练或评估结果不对时我建议按输入数据、网络输出、损失计算、指标计算这个顺序排查不要一上来就改网络结构。先用一张固定 patch 在 CPU 上跑一次 forward确认输入输出 shape 和值域再检查 loss 曲线是否符合预期走势然后是输出图像是否有结构性伪影。大多数问题都出在数据处理和归一化而不是模型结构本身。6. 进阶用真实 TEM 图像完成验证闭环合成数据的指标再漂亮都不如直接拿真实低剂量图验证一次。这里给一个我常用的闭环流程也是论文里经常用来支撑泛化性的做法。第一步是收集几组「同一视野、不同电子剂量」的配对图低剂量当成输入高剂量当成参考图。如果实验室条件有限也可以用同一区域多次扫描后平均的结果当作参考但要注意漂移问题必须先用互相关做对齐不能直接逐像素计算 PSNR。对齐这一步在电镜数据里绕不开除非两张图本身没有任何平移旋转否则指标会被运动伪影毁掉。用 scipy 的 ndimage 做刚体对齐通常就够用。第二步是构造评估指标组合。除 PSNR 和 SSIM 外我建议额外看一个频域指标对输入、输出、参考图分别做二维 FFT比较它们径向平均功率谱的高频段。TEM 去噪最怕的就是把真实高频结构也当噪声滤掉而功率谱可以直观反映这一点如果输出图的高频段相对输入没有明显塌陷说明网络保留细节的能力是可信的。第三步是可视化技巧输出图像和参考图对比时统一使用 2% 到 98% 分位数做线性拉伸不要各自用各自的 min-max否则对比图看起来亮暗不一致容易误判。如果真实配对数不够另一个可行方案是单独用好图像上的盲噪声估计来做评估找一块均匀的真空区域统计去噪前后的噪声标准差变化再结合晶格区域的主观细节判断。这个做法不严谨但能快速给出方向性参考。我在第一次跑这类项目时就因为没有做对齐直接拿配对图算了 PSNR结果指标低到以为模型失败了折腾了几天才发现是两帧之间有平移做完对齐后指标立刻正常。所以任何评判都要建立在数据本身可信的基础上先让数据干净再谈模型好坏。希望帮到你。本文还有配套的精品资源点击获取