您现在的位置是:首页 >技术交流 >SeaFormer:轻量高效的注意力模块网站首页技术交流
SeaFormer:轻量高效的注意力模块
文章目录
摘要
论文:https://openreview.net/pdf?id=-qg8MQNrxZw
代码:https://github.com/fudan-zvg/SeaFormer
自从引入视觉TransFormer以来,许多计算机视觉任务(例如,语义分割)的前景一直由cnn压倒性地主导,最近发生了重大变革。然而,计算成本和内存需求使得这些方法不适合移动设备,特别是对于高分辨率的逐像素语义分割任务。本文介绍了一种新的用于移动语义分割的压缩增强轴向TransFormer(SeaFormer)方法。具体来说,我们使用压缩轴和细节增强的方法设计了一个通用的注意力块。它可以进一步用于创建一系列具有卓越成本效益的骨干体系结构。再加上一个轻分割头,我们在基于arm的移动设备上在ADE20K和cityscape数据集上实现了分割精度和延迟之间的最佳权衡。关键的是,我们以更好的性能和更低的延迟击败了移动友好型竞争对手和基于transformer的对手,而且没有花哨的东西。除了语义分割之外,我们进一步将提出的SeaFormer架构应用于图像分类问题,展示了作为多功能移动友好骨干网的潜力。我们的代码和模型可以在https://github.com/fudan-zvg/SeaFormer上公开获得。
1、简介
语义分割是计算机视觉中的一个基本问题,其目的是为图像中的每个像素分配一个语义类标签。传统的方法依靠叠加局部卷积核Long et al.(2015)来感知图像的长程结构信息。
自从视觉TransFormer Dosovitskiy等人(2021)引入以来,语义分割领域发生了重大变革。基于TransFormer的方法郑等(2021);Xie等人(2021)已经显著地展示了全局上下文建模的能力。然而,Transformer的计算成本和内存需求使得这些方法不适用于移动设备,尤其是高分辨率图像输入。
遵循高效操作的传统智慧,基于本地/窗口的注意力(Luong et al., 2015);Liu et al. (2021);Huang等人(2021a);Yuan等人(2021),Axial attention Huang等人(2019b);Ho et al. (2019);Wang等人(2020a),动态图消息传递Zhang等人(2020;2022b)和一些轻量级的注意机制,Hou et al. (2020);Li et al. (2021b;c;2020);Liu et al. (2018);Shen等人(2021);Xu et al. (2021);曹等(2019);Woo等人(2018);Wang等人(2020b);Choromanski et al. (2021);Chen等人(2017);介绍了Mehta和Rastegari (2022a)。
然而,由于高分辨率输入的高延迟,这些进步仍然不足以满足移动设备的设计要求和限制(见图1)。最近,人们对构建基于transformer的语义分割产生了浓厚的兴趣。为了在高分辨率下降低计算成本,TopFormer Zhang等人(2022c)致力于在原始输入的1/64尺度上应用全局关注,这肯定会损害分割性能。
为了在性能无害的情况下解决像素级分割任务的高分辨率计算和移动设备的低延迟要求的难题,我们提出了一种基于适用移动端TransFormer的语义分割模型,称为压缩增强轴向变压器(SeaFormer),该模型将轴向注意力的计算复杂度从 O ( ( H + W ) H W ) mathcal{O}((H+W) H W) O((H+W)HW)降低到 O ( H W ) mathcal{O}(H W) O(HW)。在移动设备上实现卓越的精度和效率权衡,填补了移动友好型高效TransFormer的空白。
核心构建Block squeeze-enhanced Axial attention(SEA attention)试图沿着水平/垂直轴将输入特征映射压缩(池)成紧凑的列/行,并计算自注意。我们将查询、键和值连接起来,以补偿在压缩过程中牺牲的细节信息,然后将其馈送到深度卷积层中以增强局部细节。
结合轻型分割头,我们的设计(见图2)在小规模特征中使用所提出的SeaFormer层,能够在移动设备上以低延迟进行高分辨率图像语义分割。如图1所示,本文提出的SeaFormer在ADE20K数据集上以更低的延迟优于其他高效神经网络。特别是,在基于arm的移动设备上,SeaFormer-Base优于轻量级CNN对应的MobileNetV3 (41.0 vs.33.1 mIoU),延迟更低(106ms vs.126ms)。
我们做出了以下贡献:
- (i)我们引入了一种新的squeeze-enhanced Axial Transformer(SeaFormer)框架,用于移动语义分割;
- (ii)设计了一个通用的注意块,其特点是压缩轴向和细节增强;它可以用于创建一系列具有优越成本效益的骨干架构;
- (iii)我们在ADE20K和cityscape数据集上展示了最佳性能,击败了移动友好型竞争对手和基于transformer的细分模型,具有明显的优势;
- (iv)除了语义分割,我们进一步将提出的SeaFormer架构应用于图像分类问题,展示了作为多功能移动友好骨干网的潜力。
transformer和卷积的结合。卷积是相对高效的,但不适合捕捉长程依赖关系,视觉Transformer具有强大的全局感受野能力,但由于自注意力的计算而缺乏效率。为了充分利用两者的优势,MobileViT Mehta & Rastegari (2022a)、TopFormer Zhang等人(2022c)、LVT Yang等人(2022)、Mobile-Former Chen等人(2022b)、EdgeViTs Pan等人(2022a)、MobileViTv2 Mehta & Rastegari (2022b)、EdgeFormer Zhang等人(2022a)和EfficientFormer Li等人(2022)通过结合卷积和transformer构建为高效vit。Mobile-ViT, Mobile-Former, TopFormer和EfficientFormer受变压器模块的限制,在模型设计中必须权衡效率和性能。LVT、MobileViTv2和EdgeViTs以相对较高的计算代价保持模型大小较小,这也意味着较高的延迟。
轴向注意和变体。轴向注意力Huang et al. (2019b);Ho等人(2019);Wang et al. (2020a)旨在降低原始全局自注意力的计算复杂度Vaswani et al.(2017)。它一次在单个轴上计算自注意力,并堆叠一个水平和垂直轴注意力模块以获得全局感受野。Strip pooling Hou等人(2020)和Coordinate attention Hou等人(2021)使用带状形状池化窗口沿水平或垂直维度进行池化,以收集长程上下文。Kronecker注意力网络Gao等人(2020)使用并置水平和水平平均矩阵来平均输入矩阵并执行注意力操作。与axis attention相比,这些方法和其他类似的实现以相当低的计算成本提供了性能的提高。然而,它们忽略了池化/平均操作带来的局部细节的缺乏。
移动语义分割。目前主流的高效分割方法都是基于轻量级的cnn。DFANet Li et al.(2019)采用轻量级主干来减少计算成本,并添加了特征聚合模块来细化高层和低层特征。ICNet Zhao et al.(2018)设计了一种图像级联网络来加速算法,而BiSeNet Yu et al. (2018;2021)分别提出了低层细节和高层上下文信息的双流路径。fast - scnn Poudel等人(2019)分担多分支网络的计算成本,以产生运行时快速分割CNN。TopFormer Zhang等人(2022c)提出了一种结合cnn和ViT的新架构,并在移动语义分割的准确性和计算成本之间实现了良好的权衡。然而,它仍然受到全局自注意力的繁重计算负荷的限制。
3、方法
3.1、总体架构
受双分支架构的启发Yu et al. (2021);Poudel等人(2019);Hong等人(2021);Huang et al. (2021b);Chen等人(2022b),设计了一种压缩增强的轴向变压器(SeaFormer)框架。如图2所示,SeaFormer由共享STEM、上下文分支、空间分支、融合块和光分割头等部分组成。为了进行公平的比较,我们遵循TopFormer Zhang等人(2022c)设计的STEM。它由一个步长为2的常规卷积和四个MobileNet块组成,其中第一个和第三个块的步长为2。上下文分支和空间分支共享生成的特征图,这允许我们构建快速的语义分割模型。
上下文的分支。context分支旨在从feature map x s x_s xs中捕获上下文丰富的信息。如图2中的红色分支所示,上下文分支分为三个阶段。为获得更大的感受野,在应用MobileNet块进行下采样和扩展特征维度后,堆叠SeaFormer层。与作为下采样模块的标准卷积相比,MobileNet块增加了模型的表示能力,同时保持了较低的计算量和延迟。对于除了SeaFormer- large之外的变体,SeaFormer层被应用在最后两个阶段,以在精度和效率之间取得更好的权衡。对于SeaFormer- large,我们在上下文分支的每个阶段中插入SeaFormer层。为了在分割精度和推理速度之间取得良好的权衡,本文设计了一个挤压增强的轴向注意力块(SEA注意力),在下一小节中说明。
空间分支。空间分支旨在获得高分辨率的空间信息。与上下文分支相同,空间分支重用了特征映射xs。然而,早期卷积层的特征包含丰富的空间细节,但缺乏高层语义信息。因此,设计了一个融合块,将上下文分支中的特征融合到空间分支中,将高层语义信息融入到低层空间信息中。
融合块如图2所示,空间分支中的高分辨率特征图之后是1 × 1卷积和批量归一化层,以产生要融合的特征。上下文分支中的低分辨率特征图被送入1 × 1卷积层、批量归一化层、sigmoid层,并通过双线性插值上采样到高分辨率以产生语义权重。然后,将上下文分支的语义权重与空间分支的高分辨率特征进行元素乘;融合块使低层空间特征获得高层语义信息。
轻量级分割头。最后一个融合块之后的特征直接送入提出的分割头,如图2所示。为了快速推理,我们的光分割头由两个卷积层组成,后面分别是一个批归一化层,来自第一个批归一化层的特征被馈送到一个激活层。
3.2、压缩增强轴向注意力
全局注意力可以表示为:
y
o
=
∑
p
∈
G
(
o
)
softmax
p
(
q
o
⊤
k
p
)
v
p
(1)
mathbf{y}_{o}=sum_{p in mathcal{G}(o)} operatorname{softmax}_{p}left(mathbf{q}_{o}^{ op} mathbf{k}_{p}
ight) mathbf{v}_{p} ag{1}
yo=p∈G(o)∑softmaxp(qo⊤kp)vp(1)
其中
x
∈
R
H
×
W
×
C
mathbf{x} in mathbb{R}^{H imes W imes C}
x∈RH×W×C。
q
,
k
,
v
mathbf{q}, mathbf{k}, mathbf{v}
q,k,v是x的线性投影,即
q
=
W
q
x
,
k
=
W
k
x
,
v
=
W
v
x
mathbf{q}=mathbf{W}_{q} mathbf{x}, mathbf{k}=mathbf{W}_{k} mathbf{x}, mathbf{v}=mathbf{W}_{v} mathbf{x}
q=Wqx,k=Wkx,v=Wvx,其中
W
q
,
W
k
∈
R
C
q
k
×
C
,
W
v
∈
R
C
v
×
C
mathbf{W}_{q}, mathbf{W}_{k} in mathbb{R}^{C_{q k} imes C}, mathbf{W}_{v} in mathbb{R}^{C_{v} imes C}
Wq,Wk∈RCqk×C,Wv∈RCv×C是可学习权重。
G
(
o
)
mathcal{G}(o)
G(o)表示位置
o
=
(
i
,
j
)
o=(i, j)
o=(i,j)的特征图上的所有位置。将传统注意力模块应用于
H
×
W
×
C
H imes W imes C
H×W×C的特征图时,时间复杂度为
O
(
H
2
W
2
(
C
q
k
+
C
v
)
)
mathcal{O}left(H^{2} W^{2}left(C_{q k}+C_{v}
ight)
ight)
O(H2W2(Cqk+Cv)),导致效率低、延迟高。
y
o
=
∑
p
∈
N
m
×
m
(
o
)
softmax
p
(
q
o
⊤
k
p
)
v
p
(2)
mathbf{y}_{o}=sum_{p in mathcal{N}_{m imes m}(o)} operatorname{softmax}_{p}left(mathbf{q}_{o}^{ op} mathbf{k}_{p}
ight) mathbf{v}_{p} ag{2}
yo=p∈Nm×m(o)∑softmaxp(qo⊤kp)vp(2)
y o = ∑ p ∈ N 1 × W ( o ) softmax p ( q o ⊤ k p ) v p + ∑ p ∈ N H × 1 ( o ) softmax p ( q o ⊤ k p ) v p (3) mathbf{y}_{o}=sum_{p in mathcal{N}_{1 imes W}(o)} operatorname{softmax}_{p}left(mathbf{q}_{o}^{ op} mathbf{k}_{p} ight) mathbf{v}_{p}+sum_{p in mathcal{N}_{H imes 1}(o)} operatorname{softmax}_{p}left(mathbf{q}_{o}^{ op} mathbf{k}_{p} ight) mathbf{v}_{p} ag{3} yo=p∈N1×W(o)∑softmaxp(qo⊤kp)vp+p∈NH×1(o)∑softmaxp(qo⊤kp)vp(3)
为了提高效率,有一些工作使用了Liu et al. (2021);Huang等人(2019b);Ho等人(2019)计算局部区域内的自注意力。在方程2,3中展示了两个最具代表性的高效Transformer。方程2由基于窗口的注意力表示,Luong et al.(2015)成功地将时间复杂度降低为 O ( m 2 H W ( C q k + C v ) ) = O ( H W ) mathcal{O}left(m^{2} H Wleft(C_{q k}+C_{v} ight) ight)=mathcal{O}(H W) O(m2HW(Cqk+Cv))=O(HW),其中 N m × m ( O ) mathcal{N}_{m imes m}(O) Nm×m(O)表示相邻的 m × O m imes O m×O的m个位置,但失去了全局可接受性。公式3由轴向注意力Ho et al.(2019)表示,它仅将时间复杂度降低为 O ( ( H + W ) H W ( C q k + C v ) ) = O ( ( H W ) 1.5 ) mathcal{O}left((H+W) H Wleft(C_{q k}+C_{v} ight) ight)=mathcal{O}left((H W)^{1.5} ight) O((H+W)HW(Cqk+Cv))=O((HW)1.5),其中 N H × 1 ( O ) mathcal{N}_{H imes 1}(O) NH×1(O)表示O列的所有位置; N 1 × W ( o ) mathcal{N}_{1 imes W}(o) N1×W(o)表示集合o所在行的所有位置。
针对它们的缺点,本文提出了移动端友好的压缩增强轴向注意力,简洁的压缩轴向注意力用于全局语义提取,高效的基于卷积的细节增强核用于局部细节补充。
q
(
h
)
=
1
W
(
q
→
(
C
q
k
,
H
,
W
)
1
W
)
→
(
H
,
C
q
k
)
,
q
(
v
)
=
1
H
(
q
→
(
C
q
k
,
W
,
H
)
1
H
)
)
→
(
W
,
C
q
k
)
(4)
left.mathbf{q}_{(h)}=frac{1}{W}left(mathbf{q}^{
ightarrowleft(C_{q k}, H, W
ight)} mathbb{1}_{W}
ight)^{
ightarrowleft(H, C_{q k}
ight)}, quad mathbf{q}_{(v)}=frac{1}{H}left(mathbf{q}^{
ightarrowleft(C_{q k}, W, H
ight)} mathbb{1}_{H}
ight)
ight)^{
ightarrowleft(W, C_{q k}
ight)} ag{4}
q(h)=W1(q→(Cqk,H,W)1W)→(H,Cqk),q(v)=H1(q→(Cqk,W,H)1H))→(W,Cqk)(4)
压缩轴向注意力。为了实现更高效的计算,同时聚合全局信息,采用了一种更激进的策略。同样,
q
,
k
,
v
mathbf{q}, mathbf{k}, mathbf{v}
q,k,v首先由x得到
W
q
(
s
)
,
W
k
(
s
)
∈
R
C
q
k
×
C
,
W
v
(
s
)
∈
R
C
v
×
C
mathbf{W}_{q}^{(s)}, mathbf{W}_{k}^{(s)} in mathbb{R}^{C_{q k} imes C}, mathbf{W}_{v}^{(s)} in mathbb{R}^{C_{v} imes C}
Wq(s),Wk(s)∈RCqk×C,Wv(s)∈RCv×C。根据公式4,首先通过在水平方向上取查询特征图的平均值来实现水平挤压。同样的,右边显示了垂直方向上的垂直挤压。
z
→
(
⋅
)
mathrm{z}^{
ightarrow(cdot)}
z→(⋅)意味着对张量z的维数进行给定的置换,
1
m
∈
R
m
mathbb{1}_{m} in mathbb{R}^{m}
1m∈Rm是一个所有元素都等于1的向量。对
q
mathbf{q}
q的挤压操作同样在
k
mathbf{k}
k和
v
mathbf{v}
v上重复,因此我们最终得到
q
(
h
)
,
k
(
h
)
,
v
(
h
)
∈
R
H
×
C
q
k
,
q
(
v
)
,
k
(
v
)
,
v
(
v
)
∈
R
W
×
C
q
k
mathbf{q}_{(h)}, mathbf{k}_{(h)}, mathbf{v}_{(h)} in mathbb{R}^{H imes C_{q k}} , mathbf{q}_{(v)}, mathbf{k}_{(v)}, mathbf{v}_{(v)} in mathbb{R}^{W imes C_{q k}}
q(h),k(h),v(h)∈RH×Cqk,q(v),k(v),v(v)∈RW×Cqk。压缩操作将全局信息保留到单个轴上,从而大大减轻了如下式5所示的全局语义提取。
y
(
i
,
j
)
=
∑
p
=
1
H
softmax
p
(
q
(
h
)
i
⊤
k
(
h
)
p
)
v
(
h
)
p
+
∑
p
=
1
W
softmax
p
(
q
(
v
)
j
⊤
k
(
v
)
p
)
v
(
v
)
p
(5)
mathbf{y}_{(i, j)}=sum_{p=1}^{H} operatorname{softmax}_{p}left(mathbf{q}_{(h) i}^{ op} mathbf{k}_{(h) p}
ight) mathbf{v}_{(h) p}+sum_{p=1}^{W} operatorname{softmax}_{p}left(mathbf{q}_{(v) j}^{ op} mathbf{k}_{(v) p}
ight) mathbf{v}_{(v) p} ag{5}
y(i,j)=p=1∑Hsoftmaxp(q(h)i⊤k(h)p)v(h)p+p=1∑Wsoftmaxp(q(v)j⊤k(v)p)v(v)p(5)
特征映射的每个位置仅在两个压缩轴向特征上传播信息。虽然与公式3相比,它没有明显的计算减少,但可以通过最有效的广播操作简单地实现公式5的重复。详细信息如图3所示。压缩
q
,
k
,
v
mathbf{q}, mathbf{k}, mathbf{v}
q,k,v的时间复杂度为
O
(
(
H
+
W
)
(
2
C
q
k
+
C
v
)
)
mathcal{O}left((H+W)left(2 C_{q k}+C_{v}
ight)
ight)
O((H+W)(2Cqk+Cv)),注意力操作耗时
O
(
(
H
2
+
W
2
)
(
C
q
k
+
C
v
)
)
mathcal{O}left(left(H^{2}+W^{2}
ight)left(C_{q k}+C_{v}
ight)
ight)
O((H2+W2)(Cqk+Cv))。因此,我们的压缩轴向注意力成功地将时间复杂度降低到
O
(
H
W
)
mathcal{O}(H W)
O(HW)。
轴位编码。方程4不具有位置感知,不包含feature map的位置信息。因此,我们提出压缩轴位嵌入来压缩轴位注意力。对于压缩轴向注意力,我们通过引入位置嵌入
q
(
h
)
mathbf{q}_{(h)}
q(h),
k
(
h
)
mathbf{k}_{(h)}
k(h),使
q
(
h
)
mathbf{q}_{(h)}
q(h)和
k
(
h
)
mathbf{k}_{(h)}
k(h)都知道它们在挤压轴向特征中的位置,
r
(
h
)
q
,
r
(
h
)
k
∈
R
H
×
C
q
k
mathbf{r}_{(h)}^{q}, mathbf{r}_{(h)}^{k} in mathbb{R}^{H imes C_{q k}}
r(h)q,r(h)k∈RH×Cqk是由可学习参数
B
(
h
)
q
,
B
(
h
)
k
∈
R
L
×
C
q
k
mathbf{B}_{(h)}^{q}, mathbf{B}_{(h)}^{k} in mathbb{R}^{L imes C_{q k}}
B(h)q,B(h)k∈RL×Cqk线性插值而来。L是常数。同样,将
r
(
v
)
q
,
r
(
v
)
k
∈
R
W
×
C
q
k
mathbf{r}_{(v)}^{q}, mathbf{r}_{(v)}^{k} in mathbb{R}^{W imes C_{q k}}
r(v)q,r(v)k∈RW×Cqk应用于
q
(
v
)
,
k
(
v
)
mathbf{q}_{(v)}, mathbf{k}_{(v)}
q(v),k(v)。因此,位置感知挤压轴向注意可表示为式6。
y
(
i
,
j
)
=
∑
p
=
1
H
softmax
p
(
(
q
(
h
)
i
+
r
(
h
)
i
q
)
⊤
(
k
(
h
)
p
+
r
(
h
)
p
k
)
)
v
(
h
)
p
+
∑
p
=
1
W
softmax
p
(
(
q
(
v
)
j
+
r
(
v
)
j
q
)
⊤
(
k
(
v
)
p
+
r
(
v
)
p
k
)
)
v
(
v
)
p
(6)
egin{aligned} mathbf{y}_{(i, j)} & =sum_{p=1}^{H} operatorname{softmax}_{p}left(left(mathbf{q}_{(h) i}+mathbf{r}_{(h) i}^{q}
ight)^{ op}left(mathbf{k}_{(h) p}+mathbf{r}_{(h) p}^{k}
ight)
ight) mathbf{v}_{(h) p} \ & +sum_{p=1}^{W} operatorname{softmax}_{p}left(left(mathbf{q}_{(v) j}+mathbf{r}_{(v) j}^{q}
ight)^{ op}left(mathbf{k}_{(v) p}+mathbf{r}_{(v) p}^{k}
ight)
ight) mathbf{v}_{(v) p} end{aligned} ag{6}
y(i,j)=p=1∑Hsoftmaxp((q(h)i+r(h)iq)⊤(k(h)p+r(h)pk))v(h)p+p=1∑Wsoftmaxp((q(v)j+r(v)jq)⊤(k(v)p+r(v)pk))v(v)p(6)
细节增强内核。压缩操作虽然能有效地提取全局语义信息,但却牺牲了局部细节。因此,采用基于卷积的辅助核来增强空间细节。如图3的上路径所示,
q
,
k
,
v
mathbf{q}, mathbf{k}, mathbf{v}
q,k,v首先从
x
mathbf{x}
x得到另一个
W
q
(
e
)
,
W
k
(
e
)
∈
R
C
q
k
×
C
,
W
v
(
e
)
∈
R
C
v
×
C
mathbf{W}_{q}^{(e)}, mathbf{W}_{k}^{(e)} in mathbb{R}^{C_{q k} imes C}, mathbf{W}_{v}^{(e)} in mathbb{R}^{C_{v} imes C}
Wq(e),Wk(e)∈RCqk×C,Wv(e)∈RCv×C,并在通道维度上进行连接,然后传递到由3×3深度卷积和批处理归一化组成的块中。通过3×3卷积从
q
,
k
,
v
mathbf{q}, mathbf{k}, mathbf{v}
q,k,v中聚合辅助的局部细节,然后使用带激活函数的线性投影和批归一化将
(
2
C
q
k
+
C
v
)
left(2 C_{q k}+C_{v}
ight)
(2Cqk+Cv)维压缩到C并生成细节增强权值。最后,将增强特征与挤压轴向注意所给予的特征融合。实验部分将比较不同的增强方式,包括元素加法和乘法。
3
×
3
3 imes 3
3×3深度卷积的时间复杂度为
O
(
3
2
H
W
(
2
C
q
k
+
C
v
)
)
mathcal{O}left(3^{2} H Wleft(2 C_{q k}+C_{v}
ight)
ight)
O(32HW(2Cqk+Cv)),
1
×
1
1 imes 1
1×1深度卷积的时间复杂度为
O
(
H
W
C
(
2
C
q
k
+
C
v
)
)
mathcal{O}left(H W Cleft(2 C_{q k}+C_{v}
ight)
ight)
O(HWC(2Cqk+Cv))。其他操作(如激活)的时间可以省略。
架构和变体。我们介绍了四种变体,SeaFormer-Tiny, Small, Base和Large (T, S, B和L)。更多配置细节在补充资料中列出。
4、实验
在语义分割和图像分类任务上评估了所提出方法。首先,描述了实现细节,并与现有方法进行了比较。进行了一系列消融研究,以验证SeaFormer的设计。对提出的每个分量和重要的超参数进行了彻底的研究。
4.1、实验设置
4.4.1、数据集
我们在ADE20K Zhou等人(2017),CityScapes Cordts等人(2016)上进行了分割实验。设置交并比均值(mIoU)作为评价指标。将全精度模型转换为TNN贡献者(2019),并在带有单个高通Snapdragon 865处理器的基于arm的设备上测量延迟。
ADE20K数据集涵盖150个类别,包含25K张图像,分为20K/2K/3K用于训练,val和测试。CityScapes是语义分割的驱动数据集。它由19个类别的5000个精细标注的高分辨率图像组成。
4.1.2、实现细节
我们将ImageNet-1K Deng等人(2009)的预训练网络设置为骨干网络,最后一小节介绍了ImageNet-1K的训练细节。对于语义分割,标准的BatchNorm Ioffe & Szegedy(2015)层被synchronized BatchNorm取代。
训练。我们的实现基于公共代码库mmsegmentation Contributors(2020)。遵循TopFormer Zhang等人(2022c)的批量大小、训练迭代调度程序和数据增强策略进行公平比较。初始学习率为0.0005,权重衰减为0.01。采用“poly”学习率,调度因子为1.0。在推理过程中,我们设置了与TopFormer相同的调整大小和裁剪规则,以确保公平性。城市景观的比较包括全分辨率和半分辨率。对于全分辨率版本,训练图像被随机缩放,然后裁剪为固定大小的1024 × 1024。对于半分辨率版本,训练图像被调整为1024 × 512并随机缩放,裁剪大小为1024 × 512。
4.2、与SOAT比较
ADE20K。表1显示了SeaFormer和之前的高效backbone在ADE20K val集上的结果。比较包括参数、FLOPs、延迟和mIoU。如表1所示,SeaFormer优于这些方法,具有相同或更少的FLOPs和更低的延迟。与专门设计的全局注意力语义提取器TopFormer相比,SeaFormer以更低的延迟获得了更高的分割精度。SeaFormer-B的性能比MobileNetV3高出7.9% mIoU,而延迟更低(-16%)。结果表明,所提出的SeaFormer层显著提高了表示能力。
CityScapes。从表2可以看出,SeaFormer- s在计算成本和延迟更小的情况下,取得了与TopFormer-B相当甚至更好的结果,这证明SeaFormer在高分辨率场景下也可以很好地实现性能和延迟之间的权衡。
4.3、消融研究
在本节中,我们在提出的模型中删除了不同的自注意实现和一些重要的设计元素,包括我们的挤压增强轴向注意模块(SEA注意)和ADE20K数据集上的融合块。
各成分对SEA注意力的影响。我们进行了几种配置的实验,包括仅细节增强核、仅挤压轴向注意以及两者的融合。如表3所示,只有细节增强或挤压轴向注意的性能较差,使用细节增强内核增强挤压轴向注意在ADE20K上的性能提升为2.3% mIoU。结果表明,利用卷积的局部细节增强压缩轴向注意的全局语义特征,优化了Transformer块的特征提取能力。对于增强输入,upconv(x)和conv(x)之间存在明显的性能差距。我们得出结论,增加通道将显著提高性能。比较concat[qkv]和upconv(x),它们也对应于细节增强核和挤压轴注意力之间w/或w/o的卷积权值共享,我们可以发现共享权值使我们的模型在最小的性能损失下提高了推理效率(35.8 vs.35.9)。在增强模式方面,挤压轴向注意和细节增强内核的倍增特性比添加增强高出0.4% mIoU。
不同自注意模块的比较。为了消除我们的架构的影响,并证明SEA注意的有效性和推广能力,我们在Swin Transformer Liu et al.(2021)上进行了实验,用不同的注意块替换Swin Transformer中的窗口注意。为了进行公平的比较,我们设置了与Swin相同的训练协议、超参数和模型架构配置。当用CCAttention (CCNet)或DoubleAttention (A2-Nets)代替window attention时,它们的FLOPs要比SeaFormer和其他attention block低得多。考虑到我们可能无法严格地得出结论,我们将他们的Transformer块(包括MLP)的数量增加了一倍。由于ACmix具有与Swin相同的体系结构配置,我们借用了原始论文的结果。从表4可以看出,SeaFormer以更低的FLOPs和延迟优于其他注意机制。
融合块宽度的影响。为了研究融合块宽度对融合块嵌入的影响,在SeaFormer-Base上进行了不同嵌入维度的融合实验,M表示空间分支和上下文分支特征在两个融合块中映射到的通道。结果如表5所示。
4.4、图像分类
在ImageNet-1K Deng等人(2009)上进行了实验,其中包含来自1000个类别的128万张训练图像和50万张验证图像。我们采用AdamW Kingma & Ba(2014)优化器,使用余弦衰减学习率调度器进行600次迭代。batch大小为1024,初始学习率为0.064,权重衰减为2e-5。结果如表6所示。与其他高效方法相比,SeaFormer在延迟和精度之间取得了相对较好的权衡。
4.5、时延统计
我们对提出的SeaFormer-Tiny的延迟进行了统计,如图4所示,共享的茎占整个网络延迟的一半(49%)。上下文分支的延迟大约是总延迟的三分之一(34%),而空间分支的实际延迟相对较低(8%),因为与上下文分支共享早期的卷积层。我们的光照分割头(8%)也有助于构建光照模型的成功。
5、结论
本文提出了压缩增强的轴向Transformer (SeaFormer)用于移动语义分割,填补了移动友好的高效Transformer的空缺。创建了一系列SeaFormer骨干架构,并实现了成本效益。在ADE20K和Cityscapes上的优异性能和最低的延迟验证了其在基于arm的移动设备上的有效性。除了语义分割,还将所提出的SeaFormer架构进一步应用于图像分类问题,展示了作为多功能移动友好骨干的潜力。
致谢
国家自然科学基金(批准号:62106050)、临港实验室(批准号:62106050)资助。LG-QS-202202-07)、上海市自然科学基金(22ZR1407500)、腾讯科技开放研究基金(CCF-Tencent Open Research Fund);CCF-Tencent RAGR20210111)。