通过空间分组增强模块提高卷积网络的语义特征学习能力.

本文作者假设卷积神经网络提取的特征以分组形式存储不同层次和不同部分的语义子特征,并提出了一个空间组增强(SGE)模块,该模块可以通过为每个语义子特征组中的每个空间位置生成一个注意力因子来调整每个子特征的重要性,从而为每个单独的分组增强其特征表示能力。

SGE首先把特征分组,对每组特征在空间维度上与其全局平均池化特征做点积得到初始注意力图,对注意力图进行标准化后学习两个仿射参数(缩放和偏移),然后再经过sigmoid函数进行归一化,从而实现空间注意力机制。

SGE模块实现如下:

from torch import nn

class SpatialGroupEnhance(nn.Module):
    def __init__(self, groups=64):
        super(SpatialGroupEnhance, self).__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.ones(1, groups, 1, 1))
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()
        x = x.view(b * self.groups, -1, h, w) # [bg, c/g, h, w]
        # 沿通道维度加权融合
        xn = x * self.avg_pool(x) # [bg, c/g, h, w]
        xn = xn.sum(dim=1, keepdim=True) # [bg, 1, h, w]
        # 特征标准化
        t = xn.view(b * self.groups, -1) # [bg, hw]
        t = t - t.mean(dim=1, keepdim=True)
        std = t.std(dim=1, keepdim=True) + 1e-5
        t = t / std
        # 缩放和偏移
        t = t.view(b, self.groups, h, w) # [b, g, h, w]
        t = t * self.weight + self.bias # [b, g, h, w]
        # 归一化注意力
        t = t.view(b * self.groups, 1, h, w) # [bg, 1, h, w]
        x = x * self.sigmoid(t) # [bg, c/g, h, w]
        x = x.view(b, c, h, w)
        return x
		
if __name__ == "__main__":
    t = torch.ones((32, 256, 24, 24))
    sge = SpatialGroupEnhance(256)
    out = sge(t)
    print(out.shape)

由于SGE模块的注意力操作分别在各个特征分组内进行,因此额外引入的计算量几乎可以忽略不记。

作者展示了应用SGE模块前后特征图分布的变化情况。结果表明SGE模块显著改善了组内不同语义子特征的空间分布,并产生较大的统计方差,改善了不同语义区域的特征学习。