Local Memory Attention for Fast Video Semantic Segmentation,ETH Computer Vision Lab的一篇工作:
论文地址: 点这里
代码目前未放出,如有放出会及时在这里更新。
这篇工作主要围绕着如何将现有的对静态图像的语义分割网络拓展到对视频的语义分割。对于我个人来说,是非常appealing的一个工作。
论文笔记
论文提出了一个simple and general简单且通用的模块,可用于几乎所有对静态图像的语义分割网络,使其能够实现对视频的语义分割。
文章提出的网络结构Local Memory Attention Map
Networks(LMANet) 包含三个部分:一个Memory模块通过特征图buffer的形式整合过去帧的语义信息,模块存有对过去帧的丰富表示且随着时间不断更新;对当前帧(Query)的分割是通过两个步骤来实现的:通过Attention机制从Memory模块提取Query特征,然后通过attention-based的Fusion模块整合提取的特征与Query帧本身的特征。作者提到,文中提出的框架可以很容易的整合到任何Encoder-Decoder语义分割网络中。
LMANet网络结构如下图所示:
Memory结构
文中的Memory模块通过 K e y s ( K M ) Keys(\mathcal{K}_M) Keys(KM)和 V a l u e ( V M ) Value(\mathcal{V}_M) Value(VM)构建Memory。 K e y s ( K M ) Keys(\mathcal{K}_M) Keys(KM)的远小于 V a l u e ( V M ) Value(\mathcal{V}_M) Value(VM),使查询更有效。
对输入帧 I t I_t It,我们生成两对(Key,Value)特征,一对作为Memory ( K M t , V M t ) (K_M^t,V_M^t) (KMt,VMt),另一对作为Query ( K Q t , V Q t ) (K_Q^t,V_Q^t) (KQt,VQt)。 K M \mathcal{K}_M KM和 V M \mathcal{V}_M VM分别为之前帧相应值的级联(concatenation) K M = [ K M 1 , K M 2 , . . . , K M L − 1 ] ∈ R H × W × D K × ( L − 1 ) \mathcal{K}_M=[K_M^1,K_M^2,...,K_M^{L-1}] \in \mathbb{R}^{H\times W \times D_K \times (L-1)} KM=[KM1,KM2,...,KML−1]∈RH×W×DK×(L−1)和 V M = [ V M 1 , V M 2 , . . . , V M L − 1 ] ∈ R H × W × D v × ( L − 1 ) \mathcal{V}_M=[V_M^1,V_M^2,...,V_M^{L-1}] \in \mathbb{R}^{H\times W \times D_v \times (L-1)} VM=[VM1,VM2,...,VML−1]∈RH×W×Dv×(L−1)。
Keys的对应
对新输入的一帧,我们需要获取Memory中的语义信息帮助其分割,即找到对应的
K
M
K_M
KM和
K
Q
K_Q
KQ。如果计算
K
M
K_M
KM和
K
Q
K_Q
KQ所有空间位置的相似性,则需要进行如下计算:
G
(
K
M
,
K
Q
)
=
K
M
(
i
,
j
)
T
K
Q
(
k
,
l
)
G(K_M,K_Q) = K_M(i,j)^TK_Q(k,l)
G(KM,KQ)=KM(i,j)TKQ(k,l)
G
(
K
M
,
K
Q
)
G(K_M,K_Q)
G(KM,KQ)是一个4D张量,维度为
R
H
×
W
×
H
×
W
\mathbb{R}^{H\times W \times H \times W}
RH×W×H×W。
以上计算计算量大,内从占用大。可以利用一个prior,就是当前帧的指定位置的内容有较大可能会在过去帧差不多的位置找到。因此只需要读取空间相邻位置的memory。可以通过optical flow网络中常用的correlation layer来实现。这样,相似性就可以用4D张量
C
(
K
M
,
K
Q
)
∈
H
×
W
×
R
×
R
C(K_M,K_Q)\in {H\times W \times R \times R}
C(KM,KQ)∈H×W×R×R表示,
R
R
R表示给定位置相邻区域的半径。
C
(
K
M
,
K
Q
)
=
K
M
(
i
,
j
)
T
K
Q
(
i
+
k
,
j
+
l
)
C(K_M,K_Q) = K_M(i,j)^TK_Q(i+k,j+l)
C(KM,KQ)=KM(i,j)TKQ(i+k,j+l)
在这种情况下,
(
i
,
j
)
∈
{
1
,
.
.
.
,
H
}
×
{
1
,
.
.
.
,
W
}
(i,j)\in\{1,...,H\}\times\{1,...,W\}
(i,j)∈{1,...,H}×{1,...,W}且
(
k
,
l
)
∈
R
=
{
−
R
,
.
.
.
,
R
}
(k,l)\in \mathcal{R}=\{-R,...,R\}
(k,l)∈R={−R,...,R}表示搜索关联性的区域。
计算复杂度将从 O ( ( H W ) 2 ) \mathcal{O}((HW)^2) O((HW)2)降低到 O ( ( H W ) × R 2 ) \mathcal{O}((HW)\times R^2) O((HW)×R2)
内存读取
读取内存是核心步骤。读取内存由三步完成。
首先,要完成Query key和Memory中所有key的对应。
C
k
l
t
(
K
M
t
,
K
Q
t
)
=
K
M
t
(
i
,
j
)
T
K
Q
(
i
+
k
,
j
+
l
)
C_{kl}^t(K_M^t,K_Q^t) = K_M^t(i,j)^TK_Q(i+k,j+l)
Cklt(KMt,KQt)=KMt(i,j)TKQ(i+k,j+l)
然后,将特征相关性图(feature correlation maps)通过SoftMax转化为一系列的probability maps。
P
k
l
t
=
exp
(
C
k
l
t
(
i
,
j
)
)
∑
k
l
t
exp
(
C
k
l
t
(
i
,
j
)
)
P_{kl}^t = \frac{\exp(C_{kl}^t(i,j))}{\sum_{klt}\exp(C_{kl}^t(i,j))}
Pklt=∑kltexp(Cklt(i,j))exp(Cklt(i,j))
这些probability maps代表着attention权重,如下图所示。
最终,从Memory模块读取的出每个位置的
(
i
,
j
)
(i,j)
(i,j)的输出张量
V
M
~
\tilde{\mathcal{V}_M}
VM~。
V
M
~
=
∑
t
∈
S
∑
(
k
,
l
)
∈
m
a
t
h
c
a
l
R
P
k
l
t
(
i
.
j
)
V
M
t
(
i
+
k
,
j
+
l
)
\tilde{\mathcal{V}_M} = \sum_{t\in\mathcal{S}}\sum_{(k,l)\in mathcal{R}}P_{kl}^t(i.j)V_M^t(i+k,j+l)
VM~=∑t∈S∑(k,l)∈mathcalRPklt(i.j)VMt(i+k,j+l)
读取操作是对
V
M
t
\mathcal{V}_M^t
VMt的时域
S
\mathcal{S}
S和空间域
R
\mathcal{R}
R的权重加权,
V
M
t
∈
R
H
×
W
×
D
v
\mathcal{V}_M^t \in \mathbb{R}^{H\times W\times D_v}
VMt∈RH×W×Dv。
Memory合成
F
(
V
Q
,
V
M
~
)
=
σ
(
A
Q
(
V
Q
,
V
M
~
)
)
⋅
E
Q
(
V
Q
)
+
σ
(
A
M
(
V
Q
,
V
M
~
)
)
⋅
E
M
(
V
M
)
F(V_Q,\tilde{\mathcal{V_M}})=\sigma(A_Q(V_Q,\tilde{\mathcal{V_M}}))\cdot E_Q(V_Q) +\sigma(A_M(V_Q,\tilde{\mathcal{V_M}}))\cdot E_M(V_M)
F(VQ,VM~)=σ(AQ(VQ,VM~))⋅EQ(VQ)+σ(AM(VQ,VM~))⋅EM(VM)
其中,
F
(
V
Q
,
V
M
~
)
∈
R
H
×
W
×
D
F(V_Q,\tilde{\mathcal{V_M}})\in \mathbb{R}^{H\times W\times D}
F(VQ,VM~)∈RH×W×D是一个3维张量。
E
Q
,
A
Q
,
E
M
,
A
M
E_Q,A_Q,E_M,A_M
EQ,AQ,EM,AM为卷积层。
实验结果
Ablation Study 结果