【医学图像分割模块】GFNet —— 全局滤波器网络 —— 用 FFT 在频域做“全局混频“

【医学图像分割模块】GFNet —— 全局滤波器网络 —— 用 FFT 在频域做全局混频一、论文出处论文全名Global Filter Networks for Image Classification会议/年份NeurIPS 2021论文链接https://arxiv.org/abs/2107.00645官方代码https://github.com/raoyongming/GFNet二、模块图截自论文原文图注GFNet 的核心就是把 Transformer 里的自注意力换成「2D FFT → 频域逐元素乘可学习全局滤波器 → 2D IFFT」这一条支路token mixing 全程在频域完成。三、核心思想与作用一句话总括GFNet 用「频域里的可学习全局滤波器」替代自注意力让每个 token 以 O(N log N) 的代价和全图所有 token 交互。拆解成几步看变换到频域对特征图做 2D 实数 FFTtorch.fft.rfft2把空间特征搬到频率域。空间上分散的全局信息在频域里被压缩成一组频率分量。逐元素乘全局滤波器频域特征和一个形状为(H, W//21)的可学习复数张量做逐元素相乘。这一步就是 token mixing——每个频率分量被独立地缩放/旋转等价于在空间域做一次全局卷积global circular convolution。逆变换回空间域做 2D 逆 FFTtorch.fft.irfft2回到原分辨率再走后面的通道 MLP。为什么好全局感受野代价却很低。自注意力和 MLP-mixer 的 token mixing 是 O(N²)N 是 token 数GFNet 是 O(N log N)因为 FFT 本身是 log-linear 的中间那步乘法是 O(N)。高分辨率特征图上这个差距非常明显。参数极少。可学习的只有那个滤波器张量参数量是 H×(W//21)跟通道数无关比 QKV 投影轻得多。频域乘法 空间域全局卷积。这是信号处理里的卷积定理所以它天然建模的是「长程依赖」而不是局部窗口。归纳偏置介于 CNN 和 Transformer 之间。它不像卷积那样锁死局部性也不像纯注意力那样完全无先验鲁棒性和泛化在论文里表现都不错。四、在 U-Net 里的插入位置GFNet 这个子模块频域 token mixing 块最适合放在编码器深层和瓶颈层理由瓶颈层最划算。瓶颈处分辨率最低、通道最多但 token 数仍然可观全局建模收益大而且这里算力预算相对宽松O(N log N) 相比 O(N²) 的优势在高分辨率下才最突出放在瓶颈能吃到「全局依赖」的红利又不炸显存。编码器浅层慎用。浅层特征分辨率高、以边缘纹理为主全局混合收益有限反而 FFT/IFFT 的固定开销占比高。真要用建议只在 stage3/stage4 之后替换。解码器可选。解码器要恢复空间细节频域操作会引入轻微的空间模糊因为全局卷积是循环卷积边界处理不完美如果任务对边界极敏感如细小病灶分割建议只在解码器最深层放一个别每层都塞。一句话瓶颈层放 1~2 个 GFNet 块做全局建模编码器深层按需加浅层和解码器浅层保持卷积。五、复现代码PyTorch逐行中文注释说明这是简化教学版只保留 GFNet 最核心的两个组件——频域全局滤波GlobalFilter和它外面的残差块GFNetBlock。省略了论文里的多尺度滤波、LayerScale、DropPath 等训练技巧方便直接读懂结构。importtorchimporttorch.nnasnnclassGlobalFilter(nn.Module):频域全局滤波2D FFT - 逐元素乘可学习滤波器 - 2D IFFTdef__init__(self,dim,h,w):super().__init__()# 可学习的全局滤波器形状为 (dim, h, w//21)# 用复数存储实部 虚部两个实数张量# rfft2 的输出频率维度是 w//21实数输入的厄米对称性self.complex_weightnn.Parameter(torch.randn(dim,h,w//21,2)*0.02)self.hh self.wwdefforward(self,x):# x: (B, H, W, C) —— 注意这里用 channels-last方便和 rfft2 对齐B,H,W,Cx.shape# 1) 2D 实数 FFT输出复数张量 (B, H, W//21, C)xtorch.fft.rfft2(x,dim(1,2),normortho)# 2) 把可学习滤波器拼成复数real i*imagweighttorch.view_as_complex(self.complex_weight)# (C, H, W//21)# 3) 频域逐元素相乘广播到 batch 维—— 这一步就是 token mixingxx*weight# 4) 2D 逆 FFT回到空间域 (B, H, W, C)xtorch.fft.irfft2(x,s(H,W),dim(1,2),normortho)returnxclassGFNetBlock(nn.Module):一个最小的 GFNet 块全局滤波 通道 MLP 残差def__init__(self,dim,h,w,mlp_ratio4.0):super().__init__()self.filterGlobalFilter(dim,h,w)self.norm1nn.LayerNorm(dim)self.norm2nn.LayerNorm(dim)# 通道 MLP和 Transformer 里的 FFN 一个作用hiddenint(dim*mlp_ratio)self.mlpnn.Sequential(nn.Linear(dim,hidden),nn.GELU(),nn.Linear(hidden,dim),)defforward(self,x):# x: (B, C, H, W) —— 对外保持 NCHW方便塞进 U-NetB,C,H,Wx.shape# 转成 channels-last 给频域滤波用x_x.permute(0,2,3,1)# (B, H, W, C)x_self.norm1(x_)x_self.filter(x_)xxx_.permute(0,3,1,2)# 残差转回 NCHW# 通道 MLP 分支同样走残差x_x.permute(0,2,3,1)x_self.norm2(x_)x_self.mlp(x_)xxx_.permute(0,3,1,2)returnx六、插入示例几行塞进你的网络# 假设你的 U-Net 瓶颈层特征图是 (B, 512, 16, 16)bottlenecknn.Sequential(nn.Conv2d(256,512,3,padding1),nn.BatchNorm2d(512),nn.ReLU(inplaceTrue),GFNetBlock(dim512,h16,w16),# 直接插在瓶颈GFNetBlock(dim512,h16,w16),)七、实测经验与注意点h、w 必须和输入分辨率严格对齐。GlobalFilter里的滤波器形状是写死的输入尺寸变了就会广播报错。如果训练时用多尺度裁剪要么固定 crop size要么在 forward 里按实际 H、W 动态生成滤波器会牺牲一点速度。复杂度是真的低但常数不小。理论 O(N log N)但 FFT/IFFT 的常数因子比一次矩阵乘大在小分辨率如 7×7上未必比注意力快。分辨率越高越划算这也是建议放瓶颈/深层的原因。channels-last 和 NCHW 的来回 permute 有开销。教学版为了可读性做了 permute工程里如果整网都是 channels-last 可以省掉。别小看这几次 permute在浅层高分辨率下会明显拖速度。normortho建议保留。正交归一化让 FFT/IFFT 前后能量一致数值更稳也方便和论文对齐。滤波器初始化别用默认的 randn。论文用较小方差初始化教学版里乘了 0.02否则频域乘积一开始就把特征放大/缩小得离谱训练前期容易不稳。边界效应。频域乘法等价于循环卷积图像边界会被「卷」到对面。分割任务里如果病灶贴近图像边缘可能出现轻微伪影必要时配合 padding 或只在深层使用。八、完整工程以上代码已整理进即插即用模块仓库可直接 clone 取用https://github.com/CaiCy6/med-modules下一篇预告拆解另一个频域/全局建模方向的即插即用模块讲清它和 GFNet 在 token mixing 上的差异与适用场景。