TSE:通过局部空间上下文构造通道注意力.
在常用的通道注意力机制中,如SENet,通常使用全局平均池化为每个特征通道创建单个统计量描述符。本文主要研究了通道注意力机制所需的空间上下文统计量的数量,并发现通过原始特征的七行或七列有限局部上下文构造统计量足以匹配全局上下文统计量的性能。基于此提出了平铺挤压和激发(tiled squeeze-and-excite, TSE)模块,为每个通道创建多个统计量描述符,每个描述符仅基于局部上下文。
局部上下文描述符和全局上下文描述符通常是相似的,因此TSE是SE块的有效替代,可以直接在现有SE网络中使用,而无需重新训练。相比于SE块,TSE减少了数据处理流程的缓冲(buffer)要求,因此TSE对于将SE网络部署到数据流AI加速器具有重要的实际意义。例如使用TSE将EfficientDetD2中的SE缓冲量减少了$90\%$(从$50$M减少到$4.77$M),而不损失准确性。
SE和TSE的结构对比如下,TSE把全局平均池化替换成$n \times n$平均池化。更小的池化核使用更小的空间上下文信息,并为针对数据流设计的AI加速器中的元素乘法引入更小的缓冲区。
两个模块的结构如下。TSE使用有限范围的平均池化代替全局池化。池化的步长与池化核的空间维度相同,因此平铺不重叠。在缩放输入特征之前,使用最近邻插值将注意力分布广播回输入张量的维度。
def TSE(x, kernel, se_ratio):
# x: input feature map [N, C, H, W]
# kernel: tile size (Kh, Kw)
# se_ratio: SE channel reduction ratio
N, C, H, W = x.size()
# tiled squeeze
sq = nn.AvgPool2d(kernel, stride=kernel, ceil_mode=True)
# original se excitation
ex = nn.Sequential(
nn.Conv2d(C, C // se_ratio, 1),
nn.ReLU(inplace=True),
nn.Conv2d(C // se_ratio, C, 1),
nn.Sigmoid()
)
y = ex(sq(x))
# nearest neighbor interpolation
y = torch.repeat_interleave(y, kernel[0], dim=-2)[:,:,:H,:]
y = torch.repeat_interleave(y, kernel[1], dim=-1)[:,:,:,:W]
return x * y