文章目录
解锁通道与空间的协同增效:深度解析 Spatial Group Enhance (SGE) 机制
引言:从通道到空间——视觉注意力的新维度
在深度学习的飞速发展中,卷积神经网络(CNN)无疑是计算机视觉领域的主力军。其强大的特征提取能力使得机器能够在图像和视频中识别复杂的模式。然而,随着网络层数的加深和任务复杂度的提升,一个核心问题逐渐浮现:如何让模型更智能地关注图像中“重要”的信息,并抑制“不重要”的噪声?
注意力机制(Attention Mechanism)正是为解决此问题而生。它赋予了神经网络一种“聚焦”的能力,使其能够动态地为不同特征分配权重,从而增强关键信息、削弱冗余信息。早期的注意力机制,如Squeeze-and-Excitation Networks (SE-Net),主要关注通道注意力,即学习每个特征通道的重要性。它们通过全局平均池化(Global Average Pooling, GAP)来获取通道的全局上下文,然后通过简单的全连接层或1x1卷积来预测通道权重。这种方法高效且有效,但其固有的局限性在于:全局池化丢失了精细的空间位置信息。
为了弥补这一不足,研究者们开始探索空间注意力,旨在让模型关注特征图中哪些“位置”更重要。同时,混合注意力(如CBAM)将通道注意力和空间注意力结合起来,试图从两个维度共同提升特征表示。然而,这些方法往往在捕获长距离空间依赖方面效率不高,或者增加了显著的计算开销。
近年来,随着EfficientNet、ResNeXt等模型的兴起,“分组(Group)”的概念被引入到网络设计中,尤其是在卷积操作中,以提高计算效率和模型多样性。受此启发,注意力机制也开始尝试与分组结合。
正是在这样的背景下,Spatial Group Enhance (SGE) 模块被提出。SGE旨在通过一种新颖的分组式空间增强机制,克服传统注意力机制的局限性。它将特征图按通道分成多个组,并在每个组内独立地学习一个空间注意力权重。这种策略使得模型能够:
- 捕获更细粒度的空间信息: 每个通道组都可以学习自己的空间注意力,避免了全局池化带来的空间信息损失。
- 实现通道与空间的协同: 通过分组和内部的空间门控,SGE有效地实现了通道间(通过分组)和空间间(通过组内空间注意力)的协同增强。
- 保持计算高效: 相较于复杂的自注意力机制,SGE的计算量相对较小,易于集成。
本文将深入探讨Spatial Group Enhance (SGE) 的原理,详细解析其核心机制、如何利用分组来提升空间注意力的效率与效果,以及其在YOLOv8等现代骨干网络中的集成与应用。通过详细的解析,我们将揭示SGE如何巧妙地平衡了效率与性能,为深度学习模型注入更精细、更强大的“洞察力”。
I. 从通道注意力到空间注意力:挑战与机遇
在深入SGE之前,我们先简要回顾一下常见的注意力机制,以便更好地理解SGE的创新之处。
1. 通道注意力:SE-Net
- 原理: 通过全局平均池化(Squeeze)获取每个通道的全局上下文信息,然后通过两个全连接层(Excitation)学习通道权重,最后将权重应用于原始特征图(Scale)。
- 优点: 轻量、高效,能提升模型性能。
- 局限性: 丢失了空间位置信息。所有空间位置上的特征被平均,导致模型无法区分同一通道内的不同空间区域的重要性。
2. 空间注意力:CBAM中的空间注意力模块
- 原理: 通常先对特征图在通道维度上进行平均池化和最大池化,得到两个单通道的特征图,然后将它们拼接并通过一个卷积层学习空间权重,最后Sigmoid激活后应用于原始特征图。
- 优点: 能够聚焦于图像中的重要区域。
- 局限性: 往往是在通道注意力之后串联使用,或者其学习空间权重的方式较为粗糙,难以捕获复杂的、长距离的空间依赖。
3. 分组卷积与多分支结构
在ResNeXt、MobileNet等模型中,分组卷积(Group Convolution)被广泛使用。它将输入通道分成多组,每组独立进行卷积。这种方式有几个好处:
- 降低计算量: 相比标准卷积,分组卷积的FLOPs更少。
- 增加模型多样性: 每个组可以学习不同的特征表示,有助于提升模型容量。
- 作为多分支结构: 可以看作是一种特殊的“多分支”结构,每个分支处理一部分通道。
受分组卷积的启发,研究者们开始思考,能否将注意力机制也以分组的形式进行?例如,每个通道组独立地学习注意力,这样既能利用分组的效率,又能捕获更细粒度的信息。SGE正是这种思想的体现,但它更专注于分组式的空间注意力。
传统注意力机制的挑战与SGE的机遇
无论是通道注意力还是空间注意力,都面临一个核心挑战:如何在不显著增加计算和参数开销的前提下,有效地融合通道信息和精细的空间位置信息? 通道注意力过于关注“是什么”而忽略“在哪里”,空间注意力则可能因为计算方式过于简单而无法捕捉复杂空间关系。
SGE正是抓住了这个机遇。它不是简单地在全局层面上应用空间注意力,而是巧妙地将通道分组,让每个组独立地在空间维度上进行“自增强”,并通过学习到的可学习参数对这种增强进行调制,从而在保持计算效率的同时,实现通道与空间的深度协同。
II. Spatial Group Enhance (SGE) 原理深度解析
Spatial Group Enhance(SGE)模块的核心思想是将输入特征图的通道维度划分为多个组,并在每个组内独立地进行空间信息增强。这种“分而治之”的策略使得SGE能够更有效地捕获细粒度的空间依赖,同时保持计算效率。
让我们一步步分解SGE模块的内部机制。
输入: 假设我们有一个输入特征图 x
,其维度为 (b, c, h, w)
,其中 b
是批量大小,c
是通道数,h
是高度,w
是宽度。
参数: self.groups
,表示将通道 c
分成多少个组。
步骤1:通道分组与组内初始激活
-
通道分组(
x = x.view(b * self.groups, -1, h, w)
):
这是SGE的第一步。输入特征图x
的形状是(b, c, h, w)
。
通过x.view(b * self.groups, -1, h, w)
,它被重塑为(b * groups, c // groups, h, w)
。
这意味着:- 批次大小
b
和组数groups
被合并到了新的批次维度b * groups
。 - 每个组的通道数是
c // groups
。 - 这样,我们就可以将每个组视为一个独立的“迷你特征图”,后续的操作将在这些“迷你特征图”上并行进行。
- 核心意义: 这一步是SGE实现“分组式”空间增强的关键。它使得每个组的通道都可以独立地处理其空间信息,从而捕获更细粒度的、与特定通道子集相关的空间注意力。
- 批次大小
-
组内自加权(
xn = x * self.avg_pool(x)
):self.avg_pool = nn.AdaptiveAvgPool2d(1)
:这是一个全局平均池化层,它会将输入的(Any_N, Any_C, H, W)
形状的特征图池化为(Any_N, Any_C, 1, 1)
。- 当
self.avg_pool(x)
应用到重塑后的x
(形状b*groups, c//groups, h, w
) 时,它将对每个组内的c // groups
个通道独立地进行全局平均池化。其输出形状为(b * groups, c // groups, 1, 1)
。 - 然后,
x * self.avg_pool(x)
进行逐元素相乘。这相当于对每个组内的每个通道,用其自身的全局平均值来对该通道的每个像素进行加权。这是一种自门控机制,类似于SE-Net中的“Squeeze”部分,但这里是在每个组内、每个通道进行。它初步增强了每个通道组内具有较高激活值的区域。
步骤2:组内空间重要性图的提取与标准化
-
通道维度聚合(
xn = xn.sum(dim=1, keepdim=True)
):- 在步骤1中,
xn
的形状是(b * groups, c // groups, h, w)
。 xn.sum(dim=1, keepdim=True)
:对每个组内的所有通道进行求和。这意味着我们将每个组内部的c // groups
个通道的信息聚合到一个单一的特征图上。- 结果
xn
的形状变为(b * groups, 1, h, w)
。 - 核心意义: 这一步将每个组内的多通道信息压缩成一个单通道的空间重要性图。这个图反映了当前组在
(h, w)
维度上的激活强度分布,即“这个组在哪些空间位置更活跃”。
- 在步骤1中,
-
展平与标准化(
t = xn.view(b * self.groups, -1)
后续操作):t = xn.view(b * self.groups, -1)
:将(b * groups, 1, h, w)
展平为(b * groups, h * w)
。现在t
的每一行对应一个组,包含了该组在所有h * w
个空间位置上的重要性值。t = t - t.mean(dim=1, keepdim=True)
:对每个组(每一行)的h * w
个空间位置值进行均值归一化,使其均值为0。std = t.std(dim=1, keepdim=True) + 1e-5
:计算每个组的h * w
个空间位置值的标准差,并加上一个小的ε防止除零。t = t / std
:对每个组的空间重要性值进行标准差归一化(Z-score Normalization)。- 核心意义: 这一系列标准化操作旨在消除不同组之间激活强度差异的影响,确保每个组的学习都能在一个相对统一的尺度上进行。它使得每个组的空间重要性图在数值上更稳定,能够更有效地学习相对的空间优先级。
步骤3:可学习的组间调制与最终空间门控
-
重塑回分组结构(
t = t.view(b, self.groups, h, w)
):- 将标准化后的
t
从(b * groups, h * w)
重新重塑为(b, groups, h, w)
。现在,它清晰地表示了每个批次图像中每个组的空间重要性分布。
- 将标准化后的
-
可学习的组间调制(
t = t * self.weight + self.bias
):self.weight = nn.Parameter(torch.zeros(1,groups,1,1))
self.bias = nn.Parameter(torch.zeros(1,groups,1,1))
weight
和bias
是可学习的参数,形状为(1, groups, 1, 1)
。这意味着对于每个组,都有一个独立的学习到的缩放因子(weight
)和一个独立的偏置(bias
)。- 当
t
(形状(b, groups, h, w)
) 与weight
和bias
进行逐元素操作时,由于广播机制,weight
和bias
会沿着b, h, w
维度复制。 - 核心意义: 这是SGE实现“增强”的关键步骤。通过这两个可学习参数,模型能够自适应地调整每个组学习到的空间重要性图的强度和偏置。例如,如果某个组学习到的空间分布在训练中被发现特别有用,其对应的
weight
可能会增大,从而放大该组空间信息的影响力。这引入了跨组的、可学习的调制能力,使得模型能够动态地强调或抑制某些组的空间注意力。
-
应用Sigmoid门控(
x = x * self.sig(t)
):t = t.view(b * self.groups, 1, h, w)
:再次将调制后的t
重塑回(b * groups, 1, h, w)
,以便与x
进行匹配。self.sig = nn.Sigmoid()
:对t
应用Sigmoid激活函数,将值映射到(0, 1)
范围内,形成最终的空间注意力权重图。x = x * self.sig(t)
:将Sigmoid激活后的注意力权重图与原始分组后的特征x
(形状b*groups, c//groups, h, w
) 进行逐元素相乘。由于self.sig(t)
的通道维度是1,它会广播到每个组内的所有通道。- 核心意义: 这是最终的空间门控操作。Sigmoid输出的0到1之间的权重,有效地对每个组内的每个像素在所有通道上进行加权。如果某个空间位置的重要性高,其权重接近1,特征值得以保留;如果重要性低,权重接近0,特征值被抑制。
步骤4:恢复原始形状
x = x.view(b, c, h, w)
:将经过空间增强的特征从分组后的(b * groups, c // groups, h, w)
形状重塑回原始的(b, c, h, w)
形状,作为SGE模块的最终输出。
SGE的初始化(init_weights
函数)
init_weights
函数是PyTorch模块中常见的权重初始化方法。它确保在训练开始时,模型参数处于一个有利于收敛的初始状态。
nn.Conv2d
:使用Kaiming Normal初始化,适用于ReLU族激活函数。nn.BatchNorm2d
:weight
初始化为1,bias
初始化为0,确保初始阶段BN层不改变特征的均值和方差。nn.Linear
:使用小的高斯随机值初始化。- 对于SGE自身的
weight
和bias
参数: 虽然在init_weights
中没有显式地为其类型 (nn.Parameter
) 单独设置初始化,但在其定义self.weight=nn.Parameter(torch.zeros(1,groups,1,1))
和self.bias=nn.Parameter(torch.zeros(1,groups,1,1))
时,它们被初始化为零。这是一种合理的起始点,意味着在训练开始时,注意力权重初始值为0.5(sigmoid(0)=0.5
),对特征图没有立即的增强或抑制,而是让模型在训练中逐步学习。
SGE的优势总结:
- 细粒度空间注意力: 通过通道分组,每个组独立学习空间注意力图,克服了全局池化丢失空间信息的弊端,实现了更细致的空间感知。
- 通道与空间协同: 模块内部的设计使得通道分组和组内空间增强紧密结合,相互促进。
- 可学习的组间调制:
weight
和bias
参数允许模型自适应地调整每个组学习到的空间重要性,增加了模型的表达能力。 - 计算高效: 相比于复杂的自注意力机制,SGE主要通过重塑、池化和逐元素操作实现,计算量相对较小,易于集成到轻量级网络中。
- 即插即用: 作为一个独立的模块,SGE可以灵活地插入到现有CNN架构的任意位置,作为增强特征提取能力的组件。
III. 代码实现与解析
接下来,我们将详细解析您提供的 SGE.py
文件,对照原理部分,更好地理解其具体实现。
1. ultralytics\nn\SGE.py
文件解析
import numpy as np # 实际上在此模块中未使用,可能是通用导入
import torch
from torch import nn
from torch.nn import init # 用于权重初始化
class SpatialGroupEnhance(nn.Module):
def __init__(self, groups=8): # groups参数,默认8
super().__init__()
self.groups = groups
self.avg_pool = nn.AdaptiveAvgPool2d(1) # 用于组内全局平均池化
# 可学习的权重和偏置,用于调制每个组的空间重要性
# 形状为 (1, groups, 1, 1),会通过广播机制应用于每个组
self.weight = nn.Parameter(torch.zeros(1, groups, 1, 1))
self.bias = nn.Parameter(torch.zeros(1, groups, 1, 1))
self.sig = nn.Sigmoid() # Sigmoid激活函数,生成注意力权重
self.init_weights() # 调用自定义的权重初始化函数
def init_weights(self):
# 遍历模块内的所有子模块进行初始化
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, std=0.001)
if m.bias is not None:
init.constant_(m.bias, 0)
# 注意:self.weight和self.bias在__init__中被初始化为零张量,
# 不在此处通过init.XXX再次初始化。
def forward(self, x):
b, c, h, w = x.shape # 获取输入特征图的维度
# 1. 通道分组,重塑为 (b*groups, c//groups, h, w)
x_grouped = x.view(b * self.groups, -1, h, w)
# 2. 组内自加权:x_grouped * 组内全局平均值
# self.avg_pool(x_grouped) 形状 (b*groups, c//groups, 1, 1)
xn = x_grouped * self.avg_pool(x_grouped)
# 3. 组内通道聚合,得到每个组的单通道空间重要性图
# 形状 (b*groups, 1, h, w)
xn = xn.sum(dim=1, keepdim=True)
# 4. 展平并进行组内空间标准化
t = xn.view(b * self.groups, -1) # 形状 (b*groups, h*w)
t = t - t.mean(dim=1, keepdim=True) # 均值归一化
std = t.std(dim=1, keepdim=True) + 1e-5 # 标准差归一化
t = t / std
# 5. 重塑回 (b, groups, h, w),准备进行组间调制
t = t.view(b, self.groups, h, w)
# 6. 可学习的组间调制:t * weight + bias
# weight和bias会广播到所有空间维度
t = t * self.weight + self.bias
# 7. 重塑回 (b*groups, 1, h, w),并应用Sigmoid生成注意力权重
t = t.view(b * self.groups, 1, h, w)
spatial_attention_map = self.sig(t)
# 8. 将注意力权重应用到分组后的特征上 (逐元素相乘)
x_enhanced_grouped = x_grouped * spatial_attention_map
# 9. 恢复原始形状 (b, c, h, w)
out = x_enhanced_grouped.view(b, c, h, w)
return out
if __name__ == '__main__':
input_tensor = torch.randn(50, 512, 7, 7) # 模拟输入特征图
sge_module = SpatialGroupEnhance(groups=8) # 实例化SGE模块,分组为8
output_tensor = sge_module(input_tensor) # 前向传播
print("Input shape:", input_tensor.shape)
print("Output shape:", output_tensor.shape) # 打印输出形状,应与输入形状相同
代码与原理对应:
self.groups
: 控制分组的数量。self.avg_pool
: 用于每个组内部的全局平均池化。self.weight
,self.bias
: 可学习的参数,对应原理中的组间调制。self.sig
: Sigmoid激活函数,用于生成(0, 1)
范围内的注意力权重。x.view(b * self.groups, -1, h, w)
: 核心的分组操作。xn = x_grouped * self.avg_pool(x_grouped)
: 组内初步自加权。xn = xn.sum(dim=1, keepdim=True)
: 组内通道聚合为单通道空间图。t = t - t.mean(...)
,t = t / std
: 组内空间标准化。t = t * self.weight + self.bias
: 可学习的组间调制。x_enhanced_grouped = x_grouped * spatial_attention_map
: 最终的门控操作。
这个实现精确地遵循了Spatial Group Enhance论文中描述的结构和流程。
2. ultralytics\cfg\models\v8\yolov8-SpatialGroupEnhance.yaml
文件解析
您提供的 yolov8-SpatialGroupEnhance.yaml
文件是在标准YOLOv8n配置基础上,在骨干网络的末尾添加了 SpatialGroupEnhance
模块。
# ... (Parameters 和其他 backbone 部分保持不变) ...
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
- [-1, 1, SpatialGroupEnhance, []] # 10 <-- 在这里添加了SGE模块
# 这里args为空,因为groups参数有默认值8
# ... (head部分保持不变) ...
关键修改点:
在骨干网络的第10层(编号从0开始),即在 SPPF
模块之后,添加了 SpatialGroupEnhance
模块。这意味着来自 SPPF
的输出(通道数为1024)将作为 SGE
的输入,然后 SGE
的输出将传递给检测头。这种放置通常是为了在高级特征图上应用注意力,以增强检测器的性能。
关于 args: []
的重要说明:
SpatialGroupEnhance
模块的 __init__
函数只有一个参数 groups
(默认值是8)。
在YAML文件中 args: []
意味着会使用默认的 groups=8
。如果您想使用不同的分组数(例如 groups=4
),您需要在YAML中指定:
- [-1, 1, SpatialGroupEnhance, [4]] # 将分组数设置为4
3. ultralytics\nn\tasks.py
文件修改解析
为了让Ultralytics框架识别并正确地实例化 SpatialGroupEnhance
,需要进行两处修改:
修改点1:导入 SpatialGroupEnhance
类
# ultralytics\nn\tasks.py
# ... 其他导入 ...
from ultralytics.nn.modules import (AIFI, C1, C2, C3, C3TR, C5, C6, CBAM, ChannelAttention, Concat, Conv,
ConvTranspose, DWConv, ELAN, Focus, GhostBottleneck, GhostConv, ImagePooling,
LightConv, MLPBlock, SPPCSPC, SPPF, RepC3, RepConv, ResNetLayer, RTDETRDecoder,
SEGDecoder, Classify, DFL, Proto,
Detect, Segment) # Keep this line if it's already there
# 新增导入SpatialGroupEnhance
from ultralytics.nn.SGE import SpatialGroupEnhance # <--- 在这里添加
# ... 其他导入 ...
解释: 这一行代码确保Python解释器能够找到并加载 SpatialGroupEnhance
类。没有这一行,当 tasks.py
尝试根据YAML文件中的 SpatialGroupEnhance
模块名称来查找对应的类时,会报 NameError
。
修改点2:在 parse_model
函数中处理 SpatialGroupEnhance
模块
parse_model
函数是 tasks.py
的核心,它遍历YAML文件中定义的每一层,并创建对应的PyTorch模块。它有一个 elif m is SpatialGroupEnhance:
判断分支来处理 SpatialGroupEnhance
模块。
# ultralytics\nn\tasks.py
# ... 其他代码 ...
class Model(nn.Module):
# ... __init__ 和 _forward_augment 函数 ...
def parse_model(self, d, ch): # model_dict, input_channels(list)
# ... 其他模块处理 ...
elif m in frozenset({TorchVision, Index}): # TorchVision, Index 模块处理
c2 = ch[f]
args = [c2, *args]
# 添加SpatialGroupEnhance模块的处理
elif m is SpatialGroupEnhance: # <--- 在这里添加
# SpatialGroupEnhance模块的__init__只接受groups参数
# 它不需要输入通道数,所以这里我们直接使用args(如果yaml中提供了分组数)
# 或者使用默认值。
# 如果yaml中指定了groups,例如 [4],那么args就是 [4]。
# 如果yaml中是 [],那么args就是 [],会使用模块默认的groups=8。
pass # 不需要额外添加c2到args中,因为SGE的__init__不需要inp参数。
# ... 其他模块处理 ...
# ... 构建模块并添加到self.model列表 ...
重要更正: 针对 tasks.py
中 elif m is SpatialGroupEnhance:
块的修改,之前的模板给出的代码是 args = [c2, *args]
。但这对于SGE是不正确的。SGE的 __init__
函数是 __init__(self, groups=8)
,它不接受 inp
(输入通道数) 作为第一个参数。它只接受 groups
参数。
因此,正确的修改应该是:
# ultralytics\nn\tasks.py
# ... 其他代码 ...
class Model(nn.Module):
# ... __init__ 和 _forward_augment 函数 ...
def parse_model(self, d, ch): # model_dict, input_channels(list)
# ... 其他模块处理 ...
elif m in frozenset({TorchVision, Index}): # TorchVision, Index 模块处理
c2 = ch[f]
args = [c2, *args]
# 添加SpatialGroupEnhance模块的处理
elif m is SpatialGroupEnhance: # <--- 在这里添加
# SGE的__init__只接受groups参数,且有默认值。
# YOLOv8的tasks.py解析时,如果yaml中args为空,则直接使用模块的默认__init__。
# 如果yaml中args非空(例如 [4]),则args就是[4],直接传给模块。
# 所以这里无需修改args,保持其原样即可。
pass # 无需任何额外处理,args列表已经包含了groups参数(如果有的话),或者为空使用默认groups。
# ... 其他模块处理 ...
# ... 构建模块并添加到self.model列表 ...
这个 pass
语句意味着 tasks.py
在遇到 SpatialGroupEnhance
模块时,会直接使用 YAML 文件中 args
字段所提供的参数(如果提供了),或者不提供任何参数,让 SpatialGroupEnhance
模块使用其自身的默认值 (groups=8
)。这是正确的处理方式,因为 SpatialGroupEnhance
的 __init__
不像 Conv
或 CoordAtt
那样需要根据输入通道数动态推断第一个参数。
IV. 部署与训练:将SGE融入YOLOv8
一旦SGE模块的代码和YOLOv8的配置文件以及 tasks.py
修改完成,就可以开始使用带有SGE的YOLOv8模型进行训练和部署了。
训练流程概述:
-
准备环境: 确保您的Python环境安装了所有必要的依赖,包括PyTorch、Ultralytics库等。
-
文件放置:
- 将
SGE.py
文件放置在ultralytics/nn/
目录下。 - 将
yolov8-SpatialGroupEnhance.yaml
文件放置在ultralytics/cfg/models/v8/
目录下。 - 修改
ultralytics/nn/tasks.py
文件,按照前面提到的两个修改点进行编辑。
- 将
-
数据集准备: 准备您的目标检测数据集(例如COCO格式、YOLO格式等)。
-
开始训练: 使用Ultralytics提供的
train
命令,指定您的新模型配置和数据集。yolo detect train model=yolov8-SpatialGroupEnhance.yaml data=your_dataset.yaml epochs=100 imgsz=640 batch=16
model=yolov8-SpatialGroupEnhance.yaml
: 指定使用包含SGE的模型配置文件。data=your_dataset.yaml
: 指定您的数据集配置文件。epochs
,imgsz
,batch
等参数根据您的需求和硬件资源进行调整。
注意事项:
-
版本兼容性: 确保您使用的Ultralytics YOLO版本与本文档及代码示例兼容。不同版本可能在文件结构或API上存在细微差异。
-
groups
参数: 如果您对groups
参数有特定的需求,记得在yolov8-SpatialGroupEnhance.yaml
中进行修改,例如:- [-1, 1, SpatialGroupEnhance, [4]] # 设置groups为4
否则将使用SGE模块默认的
groups=8
。通道数c
必须是groups
的整数倍,否则会报错。对于YOLOv8,P5层的通道数通常是1024,这对于groups=8
(1024 % 8 == 0) 是没有问题的。 -
初始权重: 对于新的模型结构,通常建议从预训练的骨干网络权重开始训练,或者直接从YOLOv8的预训练权重开始微调。由于您修改了模型结构,直接加载原有YOLOv8预训练权重可能需要使用
--weights
参数,并注意如果存在层名不匹配,可能需要进行部分加载或忽略不匹配的层。在Ultralytics中,通常它会自动处理大部分情况,只加载匹配的权重。 -
性能评估: 训练过程中和训练结束后,密切关注mAP(Mean Average Precision)、召回率(Recall)和精度(Precision)等指标,以评估SGE是否有效提升了模型性能。也要关注FLOPs和推理速度,SGE通常是比较轻量级的。
V. 总结与展望
Spatial Group Enhance (SGE) 模块通过其独特的分组式空间注意力机制,为深度学习模型带来了以下显著优势:
- 高效的空间信息捕获: 它巧妙地将通道分组,并在每个组内独立地学习空间重要性,避免了全局池化带来的空间信息损失,同时保持了计算效率。
- 通道与空间的深度协同: SGE的结构使得通道维度和空间维度的信息能够在一个统一的框架内相互影响、相互增强。
- 自适应的组间调制: 可学习的
weight
和bias
参数赋予模型更大的灵活性,能够根据数据自适应地调整不同通道组空间注意力的重要性。 - 轻量且易于集成: 相较于复杂的自注意力或多分支结构,SGE的计算量较小,可以作为即插即用的模块方便地集成到各种CNN架构中,提升性能而不引入过大的负担。
在YOLOv8等目标检测模型中集成SGE,有望进一步提升模型对图像中细粒度空间特征的感知能力,从而在复杂场景、小目标或背景干扰较大的情况下,提供更精准的检测和定位。
未来注意力机制的发展将继续探索如何更高效、更智能地融合不同维度(通道、空间、多尺度等)的信息。SGE提供了一种有前景的思路,即通过“分组”这一简单而有效的方法,在效率和性能之间找到新的平衡。随着对模型解释性和可控性需求的增加,类似SGE这样能够提供更细粒度注意力控制的模块将变得越来越重要。SGE是通向更强大、更高效视觉智能系统道路上的一个有价值的探索。
原理与模块复现
创建ultralytics\nn\SGE.py
import numpy as np
import torch
from torch import nn
from torch.nn import init
class SpatialGroupEnhance(nn.Module):
def __init__(self, groups=8):
super().__init__()
self.groups=groups
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.weight=nn.Parameter(torch.zeros(1,groups,1,1))
self.bias=nn.Parameter(torch.zeros(1,groups,1,1))
self.sig=nn.Sigmoid()
self.init_weights()
def init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, std=0.001)
if m.bias is not None:
init.constant_(m.bias, 0)
def forward(self, x):
b, c, h,w=x.shape
x=x.view(b*self.groups,-1,h,w) #bs*g,dim//g,h,w
xn=x*self.avg_pool(x) #bs*g,dim//g,h,w
xn=xn.sum(dim=1,keepdim=True) #bs*g,1,h,w
t=xn.view(b*self.groups,-1) #bs*g,h*w
t=t-t.mean(dim=1,keepdim=True) #bs*g,h*w
std=t.std(dim=1,keepdim=True)+1e-5
t=t/std #bs*g,h*w
t=t.view(b,self.groups,h,w) #bs,g,h*w
t=t*self.weight+self.bias #bs,g,h*w
t=t.view(b*self.groups,1,h,w) #bs*g,1,h*w
x=x*self.sig(t)
x=x.view(b,c,h,w)
return x
if __name__ == '__main__':
input=torch.randn(50,512,7,7)
sge = SpatialGroupEnhance(groups=8)
output=sge(input)
print(output.shape)
创建ultralytics\cfg\models\v8\yolov8-SpatialGroupEnhance.yaml
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Ultralytics YOLOv8 object detection model with P3/8 - P5/32 outputs
# Model docs: https://docs.ultralytics.com/models/yolov8
# Task docs: https://docs.ultralytics.com/tasks/detect
# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024] # YOLOv8n summary: 129 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPS
s: [0.33, 0.50, 1024] # YOLOv8s summary: 129 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPS
m: [0.67, 0.75, 768] # YOLOv8m summary: 169 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPS
l: [1.00, 1.00, 512] # YOLOv8l summary: 209 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPS
x: [1.00, 1.25, 512] # YOLOv8x summary: 209 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPS
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
- [-1, 1, SpatialGroupEnhance, []] # 10
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 13], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)
修改ultralytics\nn\tasks.py
添加from ultralytics.nn.SGE import SpatialGroupEnhance
找到elif m in frozenset({TorchVision, Index}):在之后添加
elif m is SpatialGroupEnhance:
c2 = ch[f]
args = [c2, *args]