原文链接:https://arxiv.org/abs/1908.07678
代码:https://github.com/MendelXu/ANN.git
Non-local 是一种特别有用的语义分割技术,但也因其难以进行计算和占用GPU内存而受到批评。本文提出了Asymmetric Non-local Neural Network,其中有两个突出的组成部分:Asymmetric Pyramid Non-local Block(APNB) 和 Asymmetric Fusion Non-local Block(AFNB)。APNB利用金字塔采样模块,在不牺牲性能的前提下,极大地减少了计算和内存消耗;AFNB是由APNB演化而来的,在充分考虑了长期相关性的前提下,融合了不同层次的特征,从而大大提高了性能。
核心思路是:
只要query
分支和key
分支的输出保持相同的大小,Non-local的输出大小就保持不变。考虑到这一点,如果只从key
分支和value
分支中选取几个有代表性的点作为样本,就有可能在不牺牲性能的情况下显著降低时间复杂度。如下图所示,其中,
S
≪
N
S \ll N
S≪N。
Asymmetric Non-local Neural Network
提出了APNB和AFNB,其中APNB旨在减少Non-local的计算开销,AFNB提高了Non-local的学习能力,从而提高了分割性能。
重识Non-local
考虑输入特征图
X
∈
R
C
×
H
×
W
X \in \mathbb{R}^{C \times H \times W}
X∈RC×H×W,三个
1
×
1
1 \times 1
1×1卷积
W
ϕ
W_{\phi}
Wϕ、
W
θ
W_{\theta}
Wθ和
W
γ
W_{\gamma}
Wγ用来将
X
X
X变换到不同的嵌入式空间
ϕ
∈
R
C
′
×
H
×
W
\phi \in \mathbb{R}^{C' \times H \times W}
ϕ∈RC′×H×W、
θ
∈
R
C
′
×
H
×
W
\theta \in \mathbb{R}^{C' \times H \times W}
θ∈RC′×H×W、
γ
∈
R
C
′
×
H
×
W
\gamma \in \mathbb{R}^{C' \times H \times W}
γ∈RC′×H×W
ϕ
=
W
ϕ
(
X
)
,
θ
=
W
θ
(
X
)
,
γ
=
W
γ
(
X
)
\phi = W_{\phi}(X), \space \theta = W_{\theta}(X), \space \gamma = W_{\gamma}(X)
ϕ=Wϕ(X), θ=Wθ(X), γ=Wγ(X)
接着将他们三个
f
l
a
t
t
e
n
flatten
flatten为
C
′
×
N
C' \times N
C′×N,
N
N
N为空间中总像素数,
N
=
H
×
W
N = H \times W
N=H×W。相似度矩阵
V
∈
R
N
×
N
V \in \mathbb{R}^{N \times N}
V∈RN×N可由矩阵乘法计算
V
=
ϕ
T
×
θ
V = \phi^T \times \theta
V=ϕT×θ
接着,将
V
V
V进行归一化
V
⃗
=
f
(
V
)
\vec{V}=f(V)
V=f(V)
f
f
f有三种形式,分别为
s
o
f
t
m
a
x
softmax
softmax、
r
e
s
c
a
l
i
n
g
rescaling
rescaling和
n
o
n
e
none
none,这里选择用
s
o
f
t
m
a
x
softmax
softmax。对于
γ
\gamma
γ中的每个位置,都可以计算输出为
O
=
V
⃗
×
γ
T
O=\vec{V} \times \gamma^T
O=V×γT
其中,
O
∈
R
N
×
C
′
O \in \mathbb{R}^{N \times C'}
O∈RN×C′。通过参考Non-local的设计,最终的输出为
Y
=
W
o
(
O
T
)
+
X
o
r
Y
=
c
a
t
(
W
o
(
O
T
)
,
X
)
Y = W_o(O^T) + X \space or \space Y = cat(W_o(O^T),X)
Y=Wo(OT)+X or Y=cat(Wo(OT),X)
其中,
W
o
W_o
Wo使用
1
×
1
1 \times 1
1×1卷积,作为一个加权参数来调整Non-local的重要性,并且将通道维度从
C
′
C'
C′还原为
C
C
C。
Asymmetric Pyramid Nonlocal Block(APNB)
动机和分析
Non-local在两个矩阵相乘的时间复杂度为
O
(
C
′
N
2
)
=
O
(
C
′
H
2
W
2
)
O(C'N^2)=O(C'H^2W^2)
O(C′N2)=O(C′H2W2)。在语义分割任务中,网络的输出通常有较大的分辨率来保持细节的语义信息,
N
N
N通常很大。
一个更直接的
p
i
p
e
l
i
n
e
pipeline
pipeline如下:
R
N
×
C
′
×
R
C
′
×
N
→
R
N
×
N
×
R
N
×
C
′
→
R
N
×
C
′
\mathbb{R}^{N \times C'} \times \mathbb{R}^{C' \times N}\rightarrow \mathbb{R}^{N \times N} \times \mathbb{R}^{N \times C'} \rightarrow \mathbb{R}^{N \times C'}
RN×C′×RC′×N→RN×N×RN×C′→RN×C′
通过将
N
N
N变为
S
S
S(
S
≪
N
S \ll N
S≪N),输出尺寸保持不变
R
N
×
C
′
×
R
C
′
×
S
→
R
N
×
S
×
R
S
×
C
′
→
R
N
×
C
′
\mathbb{R}^{N \times C'} \times \mathbb{R}^{C' \times S}\rightarrow \mathbb{R}^{N \times S} \times \mathbb{R}^{S \times C'} \rightarrow \mathbb{R}^{N \times C'}
RN×C′×RC′×S→RN×S×RS×C′→RN×C′
将
N
N
N变为更小的
S
S
S相当于从图中采样几个有代表性的点,而不是考虑空间中所有的点,从而将计算量大大减少。
方法
在
θ
\theta
θ和
γ
\gamma
γ后,通过添加采样模块
P
θ
P_\theta
Pθ和
P
γ
P_\gamma
Pγ,来采样一些稀疏的锚点,得到
θ
P
∈
R
C
′
×
S
\theta_P \in \mathbb{R}^{C' \times S}
θP∈RC′×S和
γ
P
∈
R
C
′
×
S
\gamma_P \in \mathbb{R}^{C' \times S}
γP∈RC′×S,其中
S
S
S为采样锚点的个数。数学公式为
θ
P
=
P
θ
(
θ
)
,
γ
P
=
P
γ
(
γ
)
\theta_P = P_{\theta}(\theta),\space \gamma_P = P_{\gamma}(\gamma)
θP=Pθ(θ), γP=Pγ(γ)
ϕ
\phi
ϕ与锚点
θ
P
\theta_P
θP间的相似度矩阵
V
P
V_P
VP如下
V
P
=
ϕ
T
×
θ
P
V_P=\phi^T \times \theta_P
VP=ϕT×θP
其中,
V
P
V_P
VP是一个不对称的矩阵,大小为
N
×
S
N \times S
N×S。接下来与Non-local一样进行归一化,得到
V
P
⃗
\vec{V_P}
VP。
a
t
t
e
n
t
i
o
n
attention
attention输出为
O
P
=
V
P
⃗
×
γ
P
T
O_P=\vec{V_P} \times \gamma_P^T
OP=VP×γPT
最终的输出
Y
P
Y_P
YP为
Y
P
=
c
a
t
(
W
o
(
O
P
T
)
,
X
)
Y_P = cat(W_o(O_P^T),X)
YP=cat(Wo(OPT),X)
时间复杂度仅为
O
(
C
′
N
S
)
O(C'NS)
O(C′NS),比标准Non-local中的
O
(
C
′
N
2
)
O(C'N^2)
O(C′N2)小很多。
前面的工作表明,全局和多尺度对语义分割很有帮助,因此,在Non-local块中嵌入金字塔池化来增强全局表示。
这四个池化结果被平铺并连接起来,作为被采样的点。设置
n
⊆
{
1
,
3
,
6
,
8
}
n \subseteq \{1,3,6,8\}
n⊆{1,3,6,8},则锚点的总量为
S
=
110
=
∑
n
⊆
{
1
,
3
,
6
,
8
}
n
2
S = 110 = \sum_{n \subseteq \{1,3,6,8\}}n^2
S=110=n⊆{1,3,6,8}∑n2
Asymmetric Fusion Nonlocal Block(AFNB)
不同层级的特征融合对语义分割有帮助,将特征融合加入到Non-local中,得到Fusion Non-local Block。
标准的Non-local只有一个输入源,而FNB有两个输入源:一个高级特征图
X
h
∈
R
C
h
×
N
h
X_h\in \mathbb{R}^{C_h \times N_h}
Xh∈RCh×Nh和一个低级特征图
X
l
∈
R
C
l
×
N
l
X_l \in \mathbb{R}^{C_l \times N_l}
Xl∈RCl×Nl。同样的,
1
×
1
1 \times 1
1×1卷积
W
ϕ
h
W_{\phi}^h
Wϕh、
W
θ
l
W_{\theta}^l
Wθl和
W
γ
l
W_{\gamma}^l
Wγl被用来将
X
h
X_h
Xh和
X
l
X_l
Xl变换到嵌入式空间
ϕ
h
∈
R
C
′
×
N
h
\phi_h \in \mathbb{R}^{C' \times N_h}
ϕh∈RC′×Nh、
θ
l
∈
R
C
′
×
N
l
\theta_l \in \mathbb{R}^{C' \times N_l}
θl∈RC′×Nl和
γ
l
∈
R
C
′
×
N
l
\gamma_l \in \mathbb{R}^{C' \times N_l}
γl∈RC′×Nl
ϕ
h
=
W
ϕ
h
(
X
h
)
,
θ
l
=
W
θ
l
(
X
l
)
,
γ
l
=
W
γ
l
(
X
l
)
\phi_h=W_{\phi}^h(X_h),\space \theta_l=W_{\theta}^l(X_l),\space \gamma_l=W_{\gamma}^l(X_l)
ϕh=Wϕh(Xh), θl=Wθl(Xl), γl=Wγl(Xl)
ϕ
h
\phi_h
ϕh和
θ
l
\theta_l
θl间的相似度矩阵
V
F
∈
R
N
h
×
N
l
V_F \in \mathbb{R}^{N_h \times N_l}
VF∈RNh×Nl为
V
F
=
ϕ
h
T
×
θ
l
V_F=\phi_h^T \times \theta_l
VF=ϕhT×θl
同样的,将
V
F
V_F
VF进行归一化得到
V
F
⃗
∈
R
N
h
×
N
l
\vec{V_F} \in \mathbb{R}^{N_h \times N_l}
VF∈RNh×Nl。同理输出为
O
F
=
V
F
⃗
×
γ
l
T
O_F=\vec{V_F} \times \gamma_l^T
OF=VF×γlT
Y P = c a t ( W o ( O P T ) , X ) Y_P = cat(W_o(O_P^T),X) YP=cat(Wo(OPT),X)
网络结构
实验结果