YOLO11 改进 - 特征融合 | STFFM空间时间特征融合模块,强化时空互补、抑制噪声,助力小目标检测高效涨点
前言
本文介绍了面向红外小目标检测的时空特征融合模块——STFFM,用于增强复杂背景下目标与噪声、杂波的区分能力。该方法通过拼接空间特征与时间/运动特征,并结合通道注意力、空间注意力和残差增强机制,实现对关键语义通道与疑似目标区域的自适应重标定。我们将 STFFM 集成进 YOLO11 的 Neck 特征融合位置,替代原有简单 Concat 操作,使跨层特征融合从直接拼接升级为注意力引导的时空/多尺度特征交互,从而提升小目标、弱纹理目标及复杂场景下的检测鲁棒性。
文章目录: YOLO11改进大全:卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总
专栏链接: YOLO11改进专栏
介绍
复杂地面场景下的红外小目标检测仍面临巨大挑战。当目标淹没于背景中时,多帧方法能够利用时域运动信息缓解该难题。然而现有多帧方法在时域运动信息提取与时空信息融合两方面均存在缺陷,限制了其检测性能。本文提出一种全局时空红外小目标检测框架。该检测框架包含两大模块:用于计算小目标相对运动图(RMM)的相对运动模式提取(RMPE)模块,以及同时接收图像与相对运动图作为输入的全局时空特征融合网络(GST-Net)。相对运动模式提取模块在各个时间步计算稠密光流图(OFMs)并做归一化处理,通过累加操作完成小目标相对运动模式的提取。全局时空特征融合网络包含两个核心子模块:时空特征融合模块(STFFM)实现特征编码中间阶段的特征交互;全局时空依赖提取模块(GSTDEM)在深层时空语义特征图中挖掘目标与背景间的长距离依赖关系,提升模型区分目标、噪声与背景的能力。多帧数据集上的实验结果表明,本文方法能够在复杂背景下有效检测小目标,性能优于其他当前最优(SOTA)算法。相关代码开源地址:https://github.com/elvintanhust/GST-Det
文章链接
论文地址:论文地址
代码地址:代码地址
基本原理
解决的关键问题
STFFM,即 Spatial-Temporal Feature Fusion Module,主要解决空间特征与时间/运动特征融合不足的问题。原论文面向红外小目标检测,目标通常尺寸极小、对比度低,容易被地物纹理、噪声和复杂背景淹没。单纯依赖空间特征时,模型难以在强杂波背景中稳定区分目标;单纯依赖时间或运动信息时,又可能因为目标运动不明显、运动噪声强或背景扰动复杂而产生漏检和误检。因此,STFFM的核心价值在于让两类特征在浅层或中间层阶段就发生交互,使空间分支提供位置、纹理和边缘信息,时间分支提供运动差异和动态线索,通过互补与过滤减少单一模态信息不足带来的检测不稳定。
整体架构
STFFM通常嵌入双分支特征编码网络中,作为空间分支和时间分支之间的桥梁。两个分支分别提取图像空间特征与运动/时序特征,在中间层将两路特征送入STFFM。模块先对两路特征进行通道对齐,然后沿通道维度拼接,形成联合的时空特征表达。随后通过通道注意力和空间注意力对融合特征进行增强,使模型自动判断哪些通道更重要、哪些空间区域更值得关注。经过多次残差式注意力增强后,再使用1×1卷积将通道数压回目标输出通道,最后用ReLU激活输出融合特征。在YOLO中,我将它设计为可直接替代neck中Concat的位置,输入两路特征,输出单路融合特征。
技术原理
STFFM的关键技术思想是“拼接融合 + 注意力筛选 + 残差增强 + 通道压缩”。拼接操作保留两路特征的完整信息;通道注意力通过全局平均池化和最大池化提取通道重要性,增强有利于目标识别的语义响应;空间注意力通过平均图和最大图生成空间权重,突出疑似目标区域并抑制背景干扰。残差结构避免注意力筛选过强导致信息损失,使原始特征和增强特征共同参与表达。最后,1×1卷积完成通道调整,使输出形状能与YOLO后续C3k2等模块自然衔接。用于YOLO11/YOLO26时,STFFM不再只是简单把特征拼起来,而是先完成一次自适应的时空/跨层特征重标定,有助于提升小目标、弱纹理目标和复杂背景场景下的检测鲁棒性。
核心代码
classfea_add_module(nn.Module):def__init__(self,channels):super(fea_add_module,self).__init__()self.ca1=ChannelAttention(channels*2)self.ca2=ChannelAttention(channels)self.sa=SpatialAttention()self.relu=nn.ReLU(inplace=True)self.shortcut1=nn.Sequential(nn.Conv2d(channels*2,channels*2,kernel_size=1,stride=1),nn.BatchNorm2d(channels*2))self.shortcut2=nn.Sequential(nn.Conv2d(channels,channels,kernel_size=1,stride=1),nn.BatchNorm2d(channels))self.center_layer=nn.Sequential(nn.Conv2d(2*channels,channels,kernel_size=3,stride=1,padding=1),nn.BatchNorm2d(channels),nn.ReLU(inplace=True),nn.Conv2d(channels,channels,kernel_size=3,padding=1),nn.BatchNorm2d(channels))defforward(self,S,T):ST=torch.cat((S,T),dim=1)out1=self.ca1(ST)*self.sa(ST)*ST res1=self.shortcut1(ST)out1+=res1 out2=self.center_layer(out1)res2=self.shortcut2(out2)out=self.ca2(out2)*self.sa(out2)*out2 out+=res2 out=self.relu(out)returnoutYOLO11引入代码
在根目录下的ultralytics/nn/目录,新建一个featureFusion目录,然后新建一个以STFFM.py为文件名的py文件, 把代码拷贝进去。
importtorchimporttorch.nnasnnclassChannelAttention(nn.Module):def__init__(self,channels,reduction=16):super().__init__()hidden=max(channels//reduction,1)self.avg_pool=nn.AdaptiveAvgPool2d(1)self.max_pool=nn.AdaptiveMaxPool2d(1)self.mlp=nn.Sequential(nn.Conv2d(channels,hidden,1,bias=False),nn.ReLU(inplace=True),nn.Conv2d(hidden,channels,1,bias=False),)self.sigmoid=nn.Sigmoid()defforward(self,x):returnself.sigmoid(self.mlp(self.avg_pool(x))+self.mlp(self.max_pool(x)))classSpatialAttention(nn.Module):def__init__(self,kernel_size=7):super().__init__()padding=kernel_size//2self.conv=nn.Conv2d(2,1,kernel_size,padding=padding,bias=False)self.sigmoid=nn.Sigmoid()defforward(self,x):avg_map=torch.mean(x,dim=1,keepdim=True)max_map,_=torch.max(x,dim=1,keepdim=True)x=torch.cat((avg_map,max_map),dim=1)returnself.sigmoid(self.conv(x))classSTFFMUnit(nn.Module):def__init__(self,channels):super().__init__()self.ca=ChannelAttention(channels)self.sa=SpatialAttention()defforward(self,x):y=x*self.sa(x)y=y*self.ca(y)returnx+yclassSTFFM(nn.Module):""" Spatial-Temporal Feature Fusion Module. YOLO-adapted implementation: two inputs -> concat -> attention refinement -> 1x1 projection. """def__init__(self,channels):super().__init__()ifisinstance(channels,int):channels=[channels,channels]iflen(channels)!=2:raiseValueError(f"STFFM expects exactly two input feature maps, but got{len(channels)}.")out_channels=channels[1]self.proj1=nn.Sequential(nn.Conv2d(channels[0],out_channels,1,bias=False),nn.BatchNorm2d(out_channels),nn.ReLU(inplace=True),)self.proj2=nn.Sequential(nn.Conv2d(channels[1],out_channels,1,bias=False),nn.BatchNorm2d(out_channels),nn.ReLU(inplace=True),)self.refine=nn.Sequential(STFFMUnit(out_channels*2),STFFMUnit(out_channels*2),STFFMUnit(out_channels*2),)self.reduce=nn.Sequential(nn.Conv2d(out_channels*2,out_channels,1,bias=False),nn.BatchNorm2d(out_channels),nn.ReLU(inplace=True),)defforward(self,x):x1,x2=x x1=self.proj1(x1)x2=self.proj2(x2)x=torch.cat((x1,x2),dim=1)x=self.refine(x)returnself.reduce(x)注册
在ultralytics/nn/tasks.py中进行如下操作:
步骤1:
fromultralytics.nn.featureFusion.STFFMimportSTFFM步骤2
修改def parse_model(d, ch, verbose=True):
elifmisSTFFM:c2=ch[f[1]]args=[[ch[x]forxinf],*args]配置yolo11-STFFM.yaml
# Ultralytics YOLO 🚀, AGPL-3.0 license# YOLO11 object detection model with STFFM fusion blocks.nc:8scales:n:[0.50,0.25,1024]s:[0.50,0.50,1024]m:[0.50,1.00,512]l:[1.00,1.00,512]x:[1.00,1.50,512]backbone:-[-1,1,Conv,[64,3,2]]-[-1,1,Conv,[128,3,2]]-[-1,2,C3k2,[256,False,0.25]]-[-1,1,Conv,[256,3,2]]-[-1,2,C3k2,[512,False,0.25]]-[-1,1,Conv,[512,3,2]]-[-1,2,C3k2,[512,True]]-[-1,1,Conv,[1024,3,2]]-[-1,2,C3k2,[1024,True]]-[-1,1,SPPF,[1024,5]]-[-1,2,C2PSA,[1024]]head:-[-1,1,nn.Upsample,[None,2,"nearest"]]-[[-1,6],1,STFFM,[]]-[-1,2,C3k2,[512,False]]-[-1,1,nn.Upsample,[None,2,"nearest"]]-[[-1,4],1,STFFM,[]]-[-1,2,C3k2,[256,False]]-[-1,1,Conv,[256,3,2]]-[[-1,13],1,STFFM,[]]-[-1,2,C3k2,[512,False]]-[-1,1,Conv,[512,3,2]]-[[-1,10],1,STFFM,[]]-[-1,2,C3k2,[1024,True]]-[[16,19,22],1,Detect,[nc]]实验
脚本
importwarnings warnings.filterwarnings('ignore')fromultralyticsimportYOLOif__name__=='__main__':# 修改为自己的配置文件地址model=YOLO('./ultralytics/cfg/models/11/yolo11-STFFM.yaml')# 修改为自己的数据集地址model.train(data='./ultralytics/cfg/datasets/coco8.yaml',cache=False,imgsz=640,epochs=10,single_cls=False,# 是否是单类别检测batch=8,close_mosaic=10,workers=0,optimizer='SGD',amp=True,project='runs/train',name='STFFM',)