在本文中,我们将介绍语义分割知识蒸馏论文Structured Knowledge Distillation for Dense Prediction。
在论文中,考虑到稠密预测是一个结构化预测问题,本文提出将结构化知识从教师网络提取到学生网络。具体来说,我们研究了两种结构化蒸馏方案:
- 成对蒸馏,通过建立静态图来提取成对相似点。
- 整体蒸馏,运用对抗性训练来提炼整体知识的整体提炼法。
论文采用了像素级蒸馏:将分割问题视为一个单独的像素标记问题的集合,使用知识蒸馏来对齐紧凑网络产生的每个像素的类别概率,使用由教师模型产生的类概率作为软目标来训练紧凑网络,损失函数如下:
其中
q
i
s
q_i^s
qis表示由紧凑网络
S
S
S产生的第
i
i
i个像素的类概率。
q
i
t
q_i^t
qit表示由繁琐网络
T
T
T产生的第
i
i
i个像素的类概率。
K
L
(
⋅
)
KL(·)
KL(⋅)是两个概率之间的Kullback-Leibler散度,
R
=
{
1
,
2
,
…
,
W
′
×
H
′
}
R=\{1,2,…, W'× H'\}
R={1,2,…,W′×H′}表示所有像素。
1. 成对蒸馏
受到成对马尔可夫随机场框架的启发,作者提出在空间位置之间转移成对关系。作者建立了一个亲和图来表示空间成对关系,其中节点表示不同的空间位置,两个节点之间的连接表示相似性。每个节点的连接范围 α \alpha α和粒度 β \beta β来控制静态亲和图的大小。
对于每个节点,论文中只根据空间距离(Chebyshev距离)和空间局部patch中聚集的 β \beta β像素来考虑与最近的 α \alpha α个节点的相似性,以表示该节点的特征,如下图所示:
成对相似蒸馏损失函数的公式如下:
其中
a
i
j
t
a^t_{ij}
aijt和
a
i
j
s
a^s_{ij}
aijs表示教师网络
T
T
T和学生网络
S
S
S产生的第
i
i
i个节点和第
j
j
j个节点之间的相似性。
R
′
=
{
1
,
2
,
…
,
W
′
×
H
′
β
}
R'=\{1,2,…, \frac{W'× H'}{\beta} \}
R′={1,2,…,βW′×H′}表示所有节点。
在论文中,作者使用平均池将
β
×
C
β×C
β×C特征聚合到
1
×
C
1×C
1×C,两个节点之间的相似性
a
i
j
a_{ij}
aij的计算从聚合的特征
f
i
f_i
fi和
f
j
f_j
fj得到:
2. 整体蒸馏
作者将教师生成的分割图与学生网络之间的高阶关系对齐。分割图的整体嵌入计算作为表示。作者使用条件生成对抗学习来描述整体蒸馏问题。整体蒸馏的损失函数如下:
其中
E
(
⋅
)
E(·)
E(⋅)为期望算子,
D
(
⋅
)
D(·)
D(⋅)为嵌入网络,在GAN中作为判别器,将
Q
Q
Q和
I
I
I投影到一个整体的嵌入评分中。Lipschitz约束通过应用梯度惩罚来实现。
将分割图与RGB图像串联作为嵌入网络 D ( ⋅ ) D(·) D(⋅)的输入。 D ( ⋅ ) D(·) D(⋅)是一个具有5个卷积块的全卷积神经网络。每个卷积块都有ReLU和BN层,除了最后输出的卷积块。在最后三个块之间插入两个自注意模块,以捕获结构信息。在输入后插入批量归一化,将分割图与RGB通道之间的比例差进行归一化。
这种判别器能够产生一个整体的嵌入,表示输入图像和分割地图的匹配程度。作者进一步添加了一个池化层,将整体嵌入到分数中。当我们在对抗训练中使用Wasserstein距离时,判别器被训练产生较高的分数,从教师网输出分割图,而从学生的输出分割图产生较低的分数。在这个过程中对评估质量的知识进行编码,将分割图转化为判别器。对学生进行规范化训练,使其在判别器的评价下获得更高的分数,从而提高学生网络的效果。
3.优化
总体目标函数由具有像素级和结构化蒸馏项的标准多类交叉熵损失组成:
其中
λ
1
\lambda_1
λ1和
λ
2
\lambda_2
λ2分别为10和0.1。作者将目标函数相对于紧凑分割网络的参数最小化,同时将目标函数相对于判别器的参数最大化,通过以下两步迭代实现:
3.1 训练判别器
训练鉴别器等价于最小化 l h o ( S , D ) l_{ho} (S,D) lho(S,D)。
目的是对教师网络中的真实样本给予较高的嵌入分数,对学生网络中的虚假样本给予较低的嵌入分数。
3.2 训练紧凑的分割网络
给定判别器网络,目标是最小化与紧凑分割网络相关的多类交叉熵损失和蒸馏损失:
其中