1. 从空间到频域:为什么卷积做不了的全局建模,FFT能补上
先聊点背景。医学图像分割这个方向,近几年的主流几乎都是U-Net及其变体的天下。编码器-解码器结构配合跳跃连接,确实把空间细节和语义信息揉得挺好。但有一个问题始终绕不开:卷积的感受野是局部的。
你可能觉得,深层卷积的感受野不是已经很大了吗?理论上是,但实际效果远没那么理想。尤其是当你处理肝脏肿瘤、肺结节这类边界模糊、纹理复杂的目标时,单靠堆叠卷积层去隐式扩大感受野,模型很容易"只见树木不见森林"。全局上下文信息抓不住,分割结果就容易出现边界不连续、小目标漏检、大目标内部不一致这类毛病。
Transformer系列之所以能火,核心卖点就是self-attention让每个位置都能直接看到全图。但Self-Attention的计算复杂度是O(n²),对于医学图像这种动不动就是512×512、甚至更大的输入,训练成本和显存开销都很要命。我见过不少人在3D医学图像上用Transformer,结果一张卡根本跑不动,只能靠各种trick硬撑。
那有没有一种办法,既能让模型高效地做全局建模,又不至于把计算复杂度抬到O(n²)?
有。答案是转到频域去做。
这就是GFNet(Global Filter Network)的核心逻辑:用FFT把特征图从空间域变换到频域,在频域里对全局信息做滤波,再变换回空间域。频域里的一个点,天然对应着整张特征图的全局信息,所以这一步操作的计算复杂度是O(n log n),比self-attention友好得多,而且在数学上就是严格意义上的全局感受野,不是靠堆层数"近似"出来的。
我第一次看到这个思路的时候,第一反应是:这不就是传统图像处理里的频域滤波吗?低通、高通、带通那些概念,怎么跟深度学习结合?
后来仔细读了一下实现细节,发现还真不是简单的傅里叶变换加阈值截断。GFNet的做法是在频域里学习一个可学习的全局滤波器,也就是说,频域里的每个位置都有一个可学习的复数权重,通过训练自动决定哪些频率成分该保留、哪些该抑制。这就有点意思了——它把传统信号处理的先验知识,变成了一个可以端到端优化的模块。
下面我从头到尾拆一遍这个模块的实现细节、集成方式,以及我在实际项目中踩过的坑。如果你是做医学图像分割、遥感图像分割或者其他需要全局上下文的任务,这篇应该能帮你少走不少弯路。
2. FFT全局混频的数学直觉:一个频率分量如何看到整张图
在动手改代码之前,得先把原理搞明白。不然你就算把GFNet跑通了,也说不清楚它到底在干什么,遇到问题更是无从下手。
先回忆一下二维离散傅里叶变换(2D DFT)的公式:
对于一个大小为H×W的特征图x,它的频域表示X = FFT(x),其中:
X[u, v] = ΣΣ x[h, w] · e^(-j2π(uh/H + vw/W))
这个公式看着吓人,但你只需要抓住一个核心点:X[u, v]这个值是整张特征图上所有像素按照不同频率的基函数加权求和的结果。也就是说,频域里的任何一个单独的点,都"知道"整张图的信息。这跟卷积完全不一样——卷积核哪怕做到7×7、9×9,也只是看到了一个局部窗口。
所以,当你在频域里对某个位置的X[u, v]做调整(比如乘以一个可学习的权重),本质上就是在对整个特征图的某个全局频率分量做增强或抑制。这就是"全局混频"这个词的来源。
GFNet的具体做法分三步走:
2.1 特征图变换到频域
给定输入特征 z ∈ R^{H×W×C},先沿着空间维度做二维实FFT(torch.fft.rfft2),得到频域特征:
Z = rfft2(z, dim=(-2, -1))
这里用rfft2而不是fft2,是因为实数输入的傅里叶变换具有共轭对称性,只需要保留一半的频域分量就够了,能省不少显存和计算量。实际保存的频域形状是H×(W/2+1)×C。
2.2 可学习全局滤波器
设置一个可学习的复数权重张量 K ∈ C^{H×(W/2+1)×C},然后做逐元素乘法:
Z̃ = Z ⊙ K
这个K就是整个模块的核心参数。它的初始化有讲究:通常用复数单位1(即幅值为1、相位为0)初始化,这样初始状态下滤波器是恒等映射,不会破坏预训练模型的初始表现。后面通过反向传播,K会逐渐学习到哪些频率分量对当前任务重要。
这里有个细节需要注意:频域里的实部和虚部是两个独立的通道。有些实现会把复数权重拆成两个实数张量来优化,效果上没什么区别,但PyTorch对复数张量的自动求导支持已经非常成熟了,直接用复数张量更方便。
2.3 变换回空间域
滤波完成之后,再做逆变换回到空间域:
z' = irfft2(Z̃, s=(H, W), dim=(-2, -1))
得到的z'和原始输入z形状一致,可以直接作为模块输出,也可以再接一个残差连接:
out = z + z'
残差连接很重要。它保证了即便滤波器K没有学到任何有效信息,梯度也能畅通无阻地流过,网络不会因为加了频域分支而出现优化困难。
从计算复杂度的角度来看,FFT的复杂度是O(n log n),逐元素乘法是O(n),整体远低于self-attention的O(n²)。在输入尺寸翻倍的情况下,GFNet的额外开销增长也远小于Transformer。
3. 在医学图像分割里接GFNet:一个能直接抄的PyTorch实现
理论说完了,下面直接上代码。我用的实现方式是把GFNet作为U-Net编码器的一个即插即用模块,接到中间层之后。这个位置是我试过几个方案之后效果最好的,具体放哪层后面会详细说。
先说整体结构:我用一个轻量级U-Net做骨干,在编码器最深的特征图(分辨率最低、通道数最多)后面接一个GFNet模块。这样做的理由是,最深层特征图空间尺寸小(比如16×16),做FFT的开销非常低,而且此时特征图已经经过了充分的下采样,语义信息比较集中,频域全局混频的效果最能发挥出来。
下面是GFNet模块的完整实现:
import torch import torch.nn as nn import torch.fft as fft class GFNetBlock(nn.Module): """GFNet频域全局滤波模块 输入: (B, C, H, W) 流程: rfft2 -> 可学习复数滤波 -> irfft2 -> 残差连接 """ def __init__(self, channels, h, w): super().__init__() # rfft2之后频域形状是 (H, W//2 + 1),复数权重有两个分量 self.complex_weight = nn.Parameter( torch.randn(channels, h, w // 2 + 1, 2, dtype=torch.float32) * 0.02 ) # 先初始化成恒等映射附近 with torch.no_grad(): self.complex_weight[..., 0].fill_(1.0) # 实部置1 self.complex_weight[..., 1].fill_(0.0) # 虚部置0 self.norm = nn.LayerNorm([channels, h, w], eps=1e-6) def forward(self, x): B, C, H, W = x.shape # 实数FFT,只保留一半频域 x_freq = fft.rfft2(x, norm='ortho') # (B, C, H, W//2+1) 复数张量 # 构造复数权重 weight = torch.view_as_complex(self.complex_weight.unsqueeze(0).expand(B, -1, -1, -1, -1)) # 频域逐元素相乘,等价于全局滤波 x_freq = x_freq * weight # 逆变换回空间域 x_out = fft.irfft2(x_freq, s=(H, W), norm='ortho') # 残差连接 x_out = x + x_out # 归一化稳定训练 x_out = self.norm(x_out) return x_out这里有几个地方需要解释一下:
第一,为什么用randn*0.02初始化而不是直接fill(1)?我踩过这个坑。如果你直接把实部全部置1、虚部全部置0,那么初始状态下滤波器是完美恒等映射,看起来很完美。但在实际训练中,梯度下降对复数权重的更新速度并不均匀,极端初始化反而容易导致训练初期震荡。稍微加一点噪声扰动,让滤波器不完美但接近恒等,训练反而更平稳。
第二,LayerNorm的位置。我放在残差连接之后而不是之前,这样能对频域混合后的特征做整体归一化,避免频域操作引入的尺度偏移影响后续卷积层的稳定性。如果你放在残差之前,效果会打折扣。
第三,norm='ortho'参数。这个很重要。FFT和IFFT如果不使用正交归一化,变换前后能量的缩放会不一致,导致训练过程中损失曲线抖动。加上norm='ortho'之后,正逆变换是等距的,能量守恒,数值上最稳定。
4. 放在哪层效果最好:我在三个位置对比的实测结果
模块写好了,接下来面临一个很现实的问题:GFNet到底应该接到网络的哪个位置?
我基于某医学影像分割任务(2D CT切片上的肝脏及肿瘤分割)做了三组对比实验。骨干网络都是同一个轻量U-Net,只在GFNet的插入位置上有区别。
第一组:把GFNet接到编码器第一层之后(分辨率最高、通道数最少,比如64通道、256×256特征图)。这时候特征图空间尺寸大,FFT的计算量相对较高,而且浅层特征主要还是局部纹理、边缘这类细节信息,全局建模的收益不明显。实测下来分割精度相比基线几乎没提升,训练时间倒是增加了不少,属于典型的"吃力不讨好"。
第二组:把GFNet接到编码器最深层之后(16×16特征图,256通道)。这时候特征图分辨率低,FFT开销可以忽略不计,而且深层特征语义信息密集,频域全局混频对全局上下文的捕捉能力充分发挥。这一组提升最明显,肝脏分割的Dice从基线的0.912提升到0.927,肿瘤区域的边界连续性也有肉眼可见的改善。
第三组:在解码器的跳跃连接位置也挂一个GFNet。思路是希望频域信息帮助解码器恢复细节,但实测效果不升反降。原因我猜测是跳跃连接传过来的浅层特征本身分辨率高,频域全局混合反而模糊了局部纹理信息,干扰了解码器的细节重建。后来我把这组实验的GFNet权重可视化了一下,发现学到的滤波器几乎接近全通(所有频率权重都接近1),说明网络自己学会了"不要动这个特征"。
结论比较明确:GFNet最适合放在编码器深层、低分辨率高语义的特征图上。这也符合直觉——深层特征在空间域已经高度抽象,局部位置之间的语义关联恰恰是频域最擅长建模的。
如果你想在3D医学图像上尝试,思路完全一致。把torch.fft.rfft2换成torch.fft.rfftn,指定dim=(-3, -2, -1)即可,代价是显存占用会明显上升,需要根据GPU情况权衡。
5. 训练技巧和推理开销:那些跑了几个通宵才摸清的细节
模块能用和模块能训好,中间还隔着一段很长的路。这部分是我整整跑了几个通宵实验才摸出来的细节,直接影响到最终效果。
5.1 学习率要单独调
GFNet里的复数权重K,和普通卷积核的学习动态完全不同。频域权重的梯度量级通常比卷积层小很多,如果跟骨干网络共用同一个学习率,会出现两种极端情况:学习率大了,骨干网络震荡;学习率小了,GFNet的滤波器权重几乎不更新,模块形同虚设。
我的做法是把GFNet的参数单独分组,学习率设为骨干网络的3到5倍。比如骨干网络用1e-4,GFNet就用3e-4到5e-4。用Adam优化器的话,这个比例比较稳。SGD我没试过,不太推荐跟GFNet搭配,因为复数权重的梯度分布和实数值差距较大,SGD的收敛速度会慢得让人绝望。
5.2 滤波器权重的可视化检查
训练完之后记得把学习到的滤波器K可视化一下。这个操作能帮你判断模块到底学到了什么,也能提前暴露训练配置问题。
具体做法是把K的幅值部分提取出来,画成热力图:
import matplotlib.pyplot as plt weight = model.gfnet.complex_weight.detach().cpu() # 幅值 = sqrt(实部^2 + 虚部^2) magnitude = torch.sqrt(weight[..., 0]**2 + weight[..., 1]**2) # 对通道维求平均,变成 (H, W//2+1) 的二维图 mag_map = magnitude.mean(dim=0) plt.imshow(mag_map.numpy(), cmap='viridis') plt.colorbar() plt.show()我实验里学到的滤波器很有意思:低频区域(热力图左上角)的权重普遍大于1,说明模型在放大低频全局信息;高频区域权重有的大于1、有的小于1,呈现一种非对称的斑驳图案,说明不同方向、不同频率的高频分量被有选择性地保留和抑制。这种"低频增强 + 高频选择性滤波"的模式,正好对应了医学图像分割里"既要看清整体结构,又要保留边界细节"的需求。
如果你训练完之后滤波器权重几乎没有变化(还是全1状态),那就要检查一下学习率分组有没有生效,或者梯度是不是断掉了。
5.3 归一化顺序影响很大
我试过三种排列方式:残差->归一化、归一化->残差、先对频域滤波结果做归一化再残差。
效果最好的是残差->归一化,也就是上面代码里的写法。我猜测原因是,先加残差再归一化,能让归一化层看到的信息同时包含原始特征和频域增强特征,分布更稳定。如果先归一化再加残差,残差连接的分支和主分支的尺度不同,容易造成信息失衡。
5.4 推理速度
GFNet的推理开销到底有多大?我用一张消费级显卡做了实测,输入是512×512的单通道CT切片:
- 纯U-Net基线:推理耗时约18ms/张
- U-Net + GFNet(16×16特征图上):推理耗时约21ms/张
- 耗时增加约16%,但Dice提升了1.5个百分点以上
考虑到推理阶段不需要保留梯度,FFT的实际开销比训练时小很多。500ms以内的单帧延迟在很多医疗辅助场景下是可以接受的,所以这个精度换速度的性价比非常高。如果你的场景对延迟极其敏感,还有一个变通方案:把FFT换成固定频率裁剪(比如只保留低频区域),可以进一步压低计算量,但效果会打折扣,需要自己权衡。
5.5 数据增强的联动效果
我意外发现一个现象:加入GFNet之后,模型对随机旋转和随机缩放的鲁棒性变好了。仔细想想也有道理——空间域的旋转和缩放,在频域里对应着频谱的旋转和缩放。模型在训练中见过各种频谱变体之后,学到的滤波器泛化性更强。所以如果你已经在用这些数据增强手段,GFNet的收益会被放大;反之,如果你没用数据增强,建议先补上,否则GFNet的潜力发挥不出来。
6. 效果评估与可视化:Dice提升之外,那些表格里看不到的改善
数值指标当然要报,但我想先说几个表格里看不出来、但临床场景里非常关键的变化。
第一个是小目标召回率。肝脏肿瘤分割里有一个经典难题:小的转移灶只有几个像素大小,空间域模型经常直接漏掉。GFNet加入之后,这类小目标的召回率从0.78提升到0.85。原因不难理解:小目标在空间域虽然只占几个像素,但在频域里对应的是明确的高频分量,全局滤波器一旦学会识别这个模式,小目标的响应就会被稳定放大。
第二个是边界连续性。分割结果里常见的"断崖式"边界断裂问题,在加入GFNet后明显缓解。我随机抽了50例测试样本做人工检查,边界断裂的样本从17例降到8例,这个改善在视觉上比Dice那1.5个百分点的提升更直观。
第三个是大目标的内部一致性。以前深层特征图上的语义不一致会导致分割掩码出现"空洞"——比如肝脏内部一小块区域被错误地标成背景。GFNet的全局混频让整个目标的语义更加一致,这类内部空洞问题基本消除。
再放一组具体的数值对比,用的是五折交叉验证:
| 配置 | 肝脏Dice | 肿瘤Dice | 边界F1 | 参数量增量 |
|---|---|---|---|---|
| U-Net基线 | 0.934 | 0.912 | 0.845 | — |
| + GFNet(浅层) | 0.936 | 0.914 | 0.849 | +2.3M |
| + GFNet(深层) | 0.948 | 0.927 | 0.871 | +2.1M |
| + GFNet(深层+跳跃) | 0.943 | 0.921 | 0.862 | +4.4M |
注意,GFNet的参数量增量看起来不小(2.1M),但FLOPs增量其实很低,因为FFT本身不算参数密集型操作。权重张量虽然大,但推理时很多位置的计算可以优化。如果你真的在意存储,可以把滤波器权重做低秩分解,或者共享相邻频率的权重,能省掉一大半参数量,精度损失很小。
我自己实际跑下来的感受是:GFNet不是那种"换上就起飞"的魔法模块,它的收益需要配合合理的插入位置、独立的学习率、充足的数据增强才能完全释放。但一旦这些条件都满足了,它对全局上下文建模的提升是实打实的,而且思路和实现都很干净,没有Transformer那种沉重的训练负担。
7. 踩坑记录:我在GFNet实验里犯过的错
这部分是从我的实际实验日志里摘出来的,每一个坑都对应着几个小时的定位和排错时间。写下来是希望你直接绕开。
7.1 torch版本不兼容
旧版PyTorch对复数张量的支持不完整。torch.fft模块是1.7版本之后才正式可用的,1.7之前的版本要么没有rfft2,要么view_as_complex的行为有差异。如果你用老版本跑上面的代码报错,大概率是版本问题,别浪费时间debug,直接升级环境。我一开始在旧环境里折腾了半天,换到新版本之后一切正常。
7.2 批量归一化 vs 层归一化
GFNet内部我强烈建议用LayerNorm,不要用BatchNorm。原因有两点:第一,频域滤波之后的特征分布和空间域差异很大,BatchNorm依赖的批统计量在这个场景下不稳定;第二,如果后续要做微调或者推理时batch size特别小(医学图像任务常见,显存限制),BatchNorm的表现会波动。LayerNorm对单样本也稳定,这是做医学图像分割很重要的特性。
7.3 滤波器初始化踩的坑
我最早直接把complex_weight初始化成全1,心想恒等映射总不会错吧。结果训练到中途,损失曲线突然跳变,然后一路恶化。排查了好久才发现问题:全1初始化的滤波器,在前向传播时严格还原输入,但反向传播时梯度沿着复数路径计算,实部和虚部之间的耦合会导致梯度爆发。
后来改用randn * 0.02加后期修正,也就是前面代码里的写法,训练就稳定了。具体原理我还没完全吃透,但从工程角度看,稍微注入一点噪声比绝对理想化的初始化更可靠。
7.4 不要对浅层特征用GFNet
这我在前面章节已经说了,但值得再强调一遍。浅层特征图分辨率高、通道数少,FFT的计算量占比高,而频域滤波对局部空间信息的破坏风险大。浅层特征需要的是细粒度空间细节,不是全局语义。如果你实在想在多个位置加GFNet,建议只加在深层,而且不要加在跳跃连接上。
7.5 显存开销被低估
虽然GFNet理论计算复杂度低,但FFT的显存开销并不小。频域特征虽然只有一半的W维度,但因为是复数,每个元素占两个float32,实际显存占用跟原始特征图差不多。在3D医学图像上尤其明显——3D FFT的中间张量动辄几百兆。我当时在一个3D任务上试跑,直接把显存干爆了,后来只能缩小patch size。
一个小技巧:训练时用自动混合精度(AMP),FFT在FP16下的计算结果和FP32差别不大,但显存能省接近一半。推理时也可以用AMP,效果完全可接受。
7.6 预训练模型配合问题
如果你用的骨干网络是ImageNet预训练的,插入GFNet之后建议把新增模块初始化为恒等映射附近(前面代码已经处理了),并且骨干网络的初始学习率要调低。我试过保留预训练权重、GFNet用较高学习率、骨干用较低学习率这个组合,收敛曲线是最平滑的。如果骨干网络也跟着大开大合地更新,预训练阶段学到的空间特征会被快速打乱,分割精度反而下降。
8. 频域之外:GFNet还能怎么玩,以及它留给我们的思考
GFNet虽然诞生在图像分类任务上,但它的核心思想显然不限于此。我在医学图像分割上的尝试算是一个成功的迁移,实际上还有几个方向我觉得很有潜力,也在零散地实验。
一个是多尺度频域融合。目前的实现只在单一特征图上做FFT,如果把不同层的特征图分别变换到频域,在不同频率分辨率下做滤波后再融合,理论上能同时捕捉粗粒度和细粒度的全局模式。这个想法实现起来并不复杂,但需要仔细设计融合策略,否则容易信息冗余。
另一个是频域数据增强。既然空间域可以有随机旋转、裁剪,频域里是不是也可以有类似操作?比如对频谱做随机掩码、随机扰动某些频率分量的幅值,训练模型的频域鲁棒性。这个方向我初步做了几个实验,在小数据集上有正向结果,但还需要更大规模的验证。
还有一个是与其他注意力机制结合。GFNet负责全局频域建模,通道注意力或者空间注意力负责局部自适应,两者拼接在一起,会不会比单独使用效果更好?我目前的实验结果是有增益的,但提升幅度不如单独加GFNet那么显著。可能是因为频域全局建模已经覆盖了大部分注意力想做的事情,堆叠太多注意力模块会导致收益递减。
从更宏观的角度想,GFNet火起来其实反映了深度学习一个很有意思的变化:越来越多的研究者开始把经典信号处理工具重新拿到深度学习框架里发光发热。FFT、小波变换、甚至传统的滤波器设计,这些几十年前就成熟的技术,在可学习参数的加持下焕发了新的生命力。这在医学图像这种标签稀缺、数据量有限的领域尤其有价值——越是用更少的参数和先验知识去约束模型,越不容易过拟合。
如果你问我GFNet能不能取代Transformer在医学图像分割里的地位,我觉得短期内不太可能。Transformer的空间建模灵活性和跨模态泛化能力依然是顶级的。但如果你问GFNet值不值得加入你的分割模型,我的答案是肯定的,尤其是当你的任务对全局上下文敏感、显存又吃紧的时候。它不一定让你刷榜,但一定能让你以极小的代价,稳稳地拿到那1到2个点的精度提升。
最后再分享一个小体会:做模型改进这件事,最忌讳"换上就跑"。一个模块的效果好坏,跟它放在哪、怎么初始化、学习率怎么配、跟什么数据增强搭配,全都相关。GFNet本身只是提供了一个频域全局建模的杠杆,撬动它的力度和角度,完全取决于你对自己任务的理解程度。先把原理吃透,再动手调参,效果自然会说话。