简介:这份资源是围绕论文《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, dose=300.0, read_noise_sigma=2.0, clip=True): """模拟低剂量TEM图像噪声。 img: 0~1范围的干净图,单通道float64 dose: 平均电子计数,越小噪声越大,常用范围20~800 read_noise_sigma: 读出噪声标准差,单位是灰度级/255 """ img = np.asarray(img, dtype=np.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, size=img.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_size=128, stride=64): 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:y+patch_size, x:x+patch_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_ch=8): super().__init__() # mid_ch 是注意力空间的维度,设得越小越省内存,但太小时表达力受限 self.q = nn.Conv2d(in_ch, mid_ch, kernel_size=1) self.k = nn.Conv2d(in_ch, mid_ch, kernel_size=1) self.v = nn.Conv2d(in_ch, in_ch, kernel_size=1) self.proj = nn.Conv2d(in_ch, in_ch, kernel_size=1) 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 范围,这一步最容易被忽略:如果你只在 dose=800 下合成数据训练,模型遇到 dose=50 的真实低剂量图像时噪声水平超出训练分布,表现会明显退化。我的习惯是让 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 loss:
import torch import torch.nn as nn class CharbonnierLoss(nn.Module): def __init__(self, eps=1e-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,结果指标低到以为模型失败了,折腾了几天才发现是两帧之间有平移,做完对齐后指标立刻正常。所以任何评判都要建立在数据本身可信的基础上,先让数据干净,再谈模型好坏。希望帮到你。
本文还有配套的精品资源,点击获取