文章目录
- 模块来源
- 模块整体简介
- 模块结构展示
- 设计出发点
- 适用场景与效果
- 代码实现与模块对应解析
- 1. 核心代码片段
- 2. 代码-论文原理对应说明
- (1)`simam_imt`
- (2)`Down_wt`
- (3)`WDE`
- (4)`CAA`
- (5)`PSAA`
- 3. 使用逻辑
模块来源
- Paper:Context-Aware and Semantic-Guided Adaptive Filtering Network for Infrared Small Target Detection
- Code:https://github.com/LingchuanK/CSAFNet
模块整体简介
PSAA(Parallel Self-Awareness Attention)是 CSAFNet 的核心特征提取模块,用于同时建模红外小目标检测中的全局上下文与局部细节。论文指出,传统自注意力虽然能捕获长程依赖,但计算代价高,且对红外小目标所依赖的高频微弱变化不敏感;纯卷积方法又难以区分与背景高度相似的小目标。为此,PSAA采用双分支并行设计:CAA分支利用红外成像物理先验与神经元能量响应建模全局信息,WDE分支借助小波变换增强高频局部细节。二者融合后,再经残差与卷积映射输出更强判别特征,是编码器与解码器中的基础单元。
模块结构展示
- 论文位置:第 III-B 节 “Parallel Self-Awareness Attention”
- 结构图:Fig. 2,标题为 “Illustration of the proposed PSAA”
- 原生结构关系:
- 输入特征
I ∈ R^{H×W×C} - 并行进入两条分支:
- CAA branch:提取非局部/全局特征
Xg - WDE branch:提取局部/高频细节特征
Xl
- CAA branch:提取非局部/全局特征
- 两分支输出相加后,经
SE block融合 - 再与输入
I做残差连接,得到Xf - 最后通过
1×1 conv与3×3 conv两路映射后相加,得到输出Y ∈ R^{H×W×2C}
- 输入特征
- 论文公式:
- 融合与残差:式 (1)
- 输出映射:式 (2)
- CAA:式 (3)–(9)
- WDE:式 (10)–(14)
设计出发点
- 现有方法缺陷:
- 基于 Transformer/self-attention 的方法虽能建模全局关系,但计算复杂,且更偏向低频平滑背景,难以捕获红外小目标依赖的高频细微梯度变化。
- 仅依赖 CNN 的局部建模又难处理目标与背景局部外观高度相似的问题。
- 模块解决的核心问题:
- 如何在较低计算代价下,同时获得全局上下文判别能力与局部高频细节感知能力。
- 设计初衷与创新思路:
- 以并行双分支替代重型自注意力。
- CAA用红外图像“目标局部能量集中、背景能量分散”的物理先验,引入 SimAM 风格神经元能量估计,增强潜在目标区域。
- WDE用 Haar 小波提取高频分量,专门强化小目标细节。
- 二者并联后融合,实现全局语义与局部细节协同建模。
适用场景与效果
- 适用场景:
- 红外小目标检测/分割
- 目标弱、尺寸小、信噪比低、背景复杂或与目标高度相似的场景
- 需要兼顾全局上下文与细粒度边缘/纹理响应的编码器-解码器网络
- 性能表现:
- 论文在第 IV-C-3 节对 PSAA 做了消融。
- 结论上,移除CAA后,
IoU、nIoU和Pd均下降,Fa明显上升,说明全局关系建模对隐藏目标识别关键。 - 移除WDE后,检测性能也显著下降,说明高频细节缺失会导致误检和漏检增加。
- 论文明确指出:双分支联合时性能最佳,证明两者具有互补性。
- 第 IV-C-2 节还表明 CAA 在编码器中的缩放因子设置
{8, 4, 2}时总体分割性能最好。
- 局限性:
- PSAA 是面向红外小目标特性设计的,依赖“目标高频弱信号 + 背景平滑”的问题设定。
- 论文虽强调其较自注意力更高效,但未在你提供内容中单独给出 PSAA 的独立计算量表。
- 模块效果依赖分支协同,若只保留单支路,性能会明显退化。
代码实现与模块对应解析
1. 核心代码片段
classsimam_imt(torch.nn.Module):def__init__(self,e_lambda=1e-4):super(simam_imt,self).__init__()self.activaton=nn.Sigmoid()self.e_lambda=e_lambdadefforward(self,x):b,c,h,w=x.size()n=w*h-1x_minus_mu_square=(x-x.mean(dim=[2,3],keepdim=True)).pow(2)y=x_minus_mu_square/(4*(x_minus_mu_square.sum(dim=[2,3],keepdim=True)/n+self.e_lambda))+0.5returnyclassDown_wt(nn.Module):def__init__(self,in_ch):super(Down_wt,self).__init__()self.wt=DWTForward(J=1,mode='zero',wave='haar')self.conv_bn_relu=nn.Conv2d(in_ch,in_ch*2,kernel_size=1,stride=1)defforward(self,x):size=x.shape[2:]yL,yH=self.wt(x)y_HL=yH[0][:,:,0,::]y_LH=yH[0][:,:,1,::]y_HH=yH[0][:,:,2,::]x=y_HL+y_LH+y_HH x=F.interpolate(x,size=size,mode='nearest')x=self.conv_bn_relu(x)returnxclassWDE(nn.Module):def__init__(self,dim):super().__init__()self.dwconv3x3=nn.Conv2d(dim,dim,kernel_size=3,padding=3//2,groups=dim)self.dwconv3x3_2=nn.Conv2d(dim,dim,kernel_size=3,padding=1,groups=dim)self.conv_0=Down_wt(in_ch=dim)self.act=nn.GELU()self.conv_1=nn.Conv2d(dim*2,dim,1,1,0)defforward(self,x):v=self.dwconv3x3(x)attn=self.dwconv3x3_2(x)attn=self.conv_0(attn)attn=self.act(attn)attn=self.conv_1(attn)attn=torch.tanh(attn)res=attn.mul(v)returnresclassCAA(nn.Module):def__init__(self,dim=36,scale=8):super(CAA,self).__init__()self.down_scale=scale self.conv1_0=nn.Conv2d(dim,dim*2,1,1,0)self.conv1_1=nn.Conv2d(dim,dim,1,1,0)self.conv1_2=nn.Conv2d(dim,dim,1,1,0)self.alpha=nn.Parameter(torch.ones((1,dim,1,1)))self.belt=nn.Parameter(torch.zeros((1,dim,1,1)))self.gelu=nn.GELU()self.WDE=WDE(dim)self.dwconv=nn.Conv2d(dim,dim,3,1,1,groups=dim)self.simam_imt=simam_imt()self.ca=nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(dim,dim//4,kernel_size=1),nn.GELU(),nn.Conv2d(dim//4,dim,kernel_size=1),nn.Sigmoid())defforward(self,f):_,_,h,w=f.shape y,x=self.conv1_0(f).chunk(2,dim=1)e_t=self.simam_imt(x)e_l=self.dwconv(F.adaptive_avg_pool2d(x,(h//self.down_scale,w//self.down_scale)))x_caa=F.interpolate(self.gelu(self.conv1_1(e_l*self.alpha)),size=(h,w),mode='nearest')+e_t*self.belt x_l=x*x_caa y_d=self.WDE(y)out=x_l+y_d out=self.ca(out)*out out=self.conv1_2(out)returnoutclassPSAA(nn.Module):def__init__(self,in_dim,out_dim,scale):super(PSAA,self).__init__()self.LG_att=CAA(in_dim,scale)self.Conv3=nn.Sequential(nn.Conv2d(in_channels=in_dim,out_channels=out_dim,kernel_size=3,stride=1,padding=1),nn.BatchNorm2d(out_dim),nn.ReLU(inplace=True))self.Conv1=nn.Sequential(nn.Conv2d(in_dim,out_dim,kernel_size=1),nn.BatchNorm2d(out_dim),nn.ReLU(inplace=True))defforward(self,x):x_f=x+self.LG_att(x)out=self.Conv3(x_f)+self.Conv1(x_f)returnout2. 代码-论文原理对应说明
(1)simam_imt
- 功能:实现神经元能量响应图计算。
- 对应论文:CAA 中的
SimAM_E(·),即式 (5)、式 (7)–(9)。 - 关键逻辑:
x.mean(dim=[2, 3], keepdim=True):按空间维求均值,对应论文中的均值统计思想。x_minus_mu_square:计算每个位置相对均值的偏差平方。... / (4 * (... + self.e_lambda)) + 0.5:与 SimAM 的最小能量函数实现一致,输出的是能量相关响应图。
- 对应关系说明:
- 论文中指出“能量越低,神经元越有区分性,重要性可由
1/e_t*体现”。 - 代码这里没有显式写
1/e_t,而是直接构造了一个可作为响应图的y,后续在CAA中作为空间调制项使用,属于论文思想的工程化实现。
- 论文中指出“能量越低,神经元越有区分性,重要性可由
(2)Down_wt
- 功能:对输入特征做 1 层 Haar 离散小波变换,提取高频子带并回到原空间尺寸。
- 对应论文:WDE 中式 (11)–(13) 前的波形分解步骤与 Fig. 3。
- 关键逻辑:
self.wt = DWTForward(J=1, mode='zero', wave='haar'):使用 Haar 小波,和论文 Fig. 3 一致。yH[0][:, :, 0, ::] / 1 / 2:分别取三个高频子带HL/LH/HH。x = y_HL + y_LH + y_HH:对应论文式 (11) 中Fh = Fq_H + Fq_V + Fq_D。F.interpolate(..., mode='nearest'):对应论文式 (12) 的UpN(Fh)。1×1 conv:对应论文式 (13) 中上采样后再卷积映射生成感知权重前的通道整合。
- 说明:
- 论文写法是先对
Fq做 wavelet transform,再通过卷积和非线性得到Fa。 - 代码将这部分封装在
Down_wt中完成高频提取和尺寸恢复。
- 论文写法是先对
(3)WDE
- 功能:生成局部细节增强特征。
- 对应论文:第 III-B-2 节,式 (10)–(14)。
- 逐步对应:
v = self.dwconv3x3(x):对应论文中的Fs,即作为被增强的局部表示。attn = self.dwconv3x3_2(x):对应论文中的Fq,用于后续生成注意/感知权重。attn = self.conv_0(attn):进入Down_wt提取高频分量并上采样,相当于论文中Fq -> [Fq_H,Fq_V,Fq_D] -> Fh -> F'h。attn = self.act(attn)与attn = self.conv_1(attn):对应式 (13) 中Conv1 -> GELU -> Conv1的非线性映射过程,代码顺序为GELU -> 1×1 conv,属于实现细节调整。attn = torch.tanh(attn):严格对应式 (13) 的tanh(...)。res = attn.mul(v):严格对应式 (14)Xl = Fa ⊙ Fs。
- 结论:
WDE与论文 WDE 分支高度一致,核心是“高频感知权重 × 局部特征”。
(4)CAA
- 功能:在代码中承担 PSAA 主体功能,实际内部同时包含“全局分支 + 局部分支 + 融合”。
- 对应论文:
- CAA 部分对应第 III-B-1 节,式 (3)–(9)
- 但代码里的
CAA不只是论文中的 CAA branch,它还把WDE(y)、融合、通道注意力等都包进来了,因此更接近“论文 PSAA 的主体实现”
- 逐步对应:
self.conv1_0 = nn.Conv2d(dim, dim * 2, 1, 1, 0)- 将输入扩展到
2C通道。
- 将输入扩展到
y, x = self.conv1_0(f).chunk(2, dim=1)- 将特征分成两半:
x走全局感知路径,y走 WDE 路径。 - 这是一种工程实现上的并行分支构造,对应论文“CAA 与 WDE 并行”。
- 将特征分成两半:
e_t = self.simam_imt(x)- 对应式 (5) 的
e_t* = SimAM_E(I')。
- 对应式 (5) 的
e_l = self.dwconv(F.adaptive_avg_pool2d(x, (h // self.down_scale, w // self.down_scale)))- 对应式 (4):先池化提取低频结构,再深度卷积建模。
self.down_scale即论文强调的 encoder/decoder 中不同阶段使用的缩放因子。
F.interpolate(..., mode='nearest')- 对应式 (4) 中
UpN(...)。
- 对应式 (4) 中
x_caa = ... + e_t * self.belt- 对应论文式 (6) 中将低频结构信息与能量响应信息结合。
- 代码中通过可学习参数
alpha、belt控制两类信息的幅度,是论文公式的参数化实现。
x_l = x * x_caa- 对应式 (6) 中最终对特征的调制。
y_d = self.WDE(y)- 调用 WDE 分支,对应论文的局部分支
Xl。
- 调用 WDE 分支,对应论文的局部分支
out = x_l + y_d- 对应论文式 (1) 里的
Xg + Xl融合。
- 对应论文式 (1) 里的
out = self.ca(out) * out- 这里加入 SE/通道注意力风格模块,对应论文式 (1) 中的
SE(...)。
- 这里加入 SE/通道注意力风格模块,对应论文式 (1) 中的
out = self.conv1_2(out)- 融合后再做
1×1 conv调整通道。
- 融合后再做
- 重要说明:
- 论文中 PSAA 写成:
Xg + Xl -> SE -> + I -> Conv1 + Conv3 - 代码中这一过程被拆开:
CAA.forward()内部完成:Xg + Xl -> channel attention -> 1×1 convPSAA.forward()再做:x + LG_att(x)以及Conv3 + Conv1
- 因此,代码整体上与论文 PSAA 对应,但模块边界划分并不完全同名。
- 论文中 PSAA 写成:
(5)PSAA
- 功能:封装完整 PSAA 输出逻辑。
- 对应论文:式 (1)、式 (2)。
- 逐步对应:
self.LG_att = CAA(in_dim, scale)- 名字虽然叫
CAA,但其内部已经包含全局与局部双分支融合,可视作“PSAA 主体注意模块”。
- 名字虽然叫
x_f = x + self.LG_att(x)- 对应论文式 (1) 的残差连接
Xf = SE(Xg + Xl) + I。
- 对应论文式 (1) 的残差连接
self.Conv3(x_f) + self.Conv1(x_f)- 对应论文式 (2):
Y = Conv1(Xf) + Conv3(Xf)。
- 对应论文式 (2):
- 输出特征:
- 论文中默认输出是
2C通道。 - 代码里通过
out_dim自定义输出维度,例如:PSAA(filters[0], filters[1], scale=8):32 -> 64PSAA(filters[1], filters[2], scale=4):64 -> 128
- 这与论文“输出升维”的描述一致。
- 论文中默认输出是
3. 使用逻辑
- 模块输入:
- 输入为四维特征图
x ∈ R^{B×C×H×W}。
- 输入为四维特征图
- 模块内部流程:
- 先进入
CAA/LG_att:- 一部分特征经 SimAM + 低频池化深度卷积,形成全局/非局部调制;
- 另一部分特征经 Haar 小波高频增强,形成局部细节响应;
- 两支路融合并经通道注意力整合。
- 再与输入残差相加。
- 最后通过
3×3与1×1两个卷积映射求和输出。
- 先进入
- 模块输出:
- 输出增强后的特征图,空间尺寸不变,通道数由
out_dim决定。
- 输出增强后的特征图,空间尺寸不变,通道数由
- 在网络中的调用方式:
- 编码器:
self.Convtans2 = PSAA(filters[0], filters[1], scale=8)self.Convtans3 = PSAA(filters[1], filters[2], scale=4)self.Convtans4 = PSAA(filters[2], filters[3], scale=2)
- 解码器:
self.Decoder1 = PSAA(filters[3], filters[2], scale=2)self.Decoder2 = PSAA(filters[2], filters[1], scale=4)self.Decoder3 = PSAA(filters[1], filters[0], scale=8)
- 编码器:
- 与论文的一致性总结:
- 代码忠实体现了 PSAA 的核心思想:并行全局-局部建模、残差融合、双卷积输出。
- 与论文的差异主要在实现封装层面:
- 代码把论文中的 CAA 与 WDE 的部分融合过程写进了
CAA类; - 代码使用可学习系数
alpha、belt调节低频结构与能量响应的贡献; - 代码用 SE 风格通道注意力
self.ca对应论文中的融合增强模块。
- 代码把论文中的 CAA 与 WDE 的部分融合过程写进了