《Structured Knowledge Distillation for Semantic Segmentation》

摘要

本文研究了利用大网络训练小语义分割网络的知识蒸馏策略，从简单的像素点蒸馏方案开始，将蒸馏方案应用于图像分类，并分别对每个像素进行知识蒸馏。此外，本文提出将结构化知识从大网络提取到小网络。本文主要研究了两种结构化的蒸馏方案：

（1）提取成对相似性的成对蒸馏；

（2）使用GAN提取整体知识的整体蒸馏。

通过对三个数据集（Cityscapes、Camvid和ADE20K）的大量实验，证明了本文知识蒸馏方法的有效性。

解决的问题

目前，最先进的语义分割模型如DeepLab, PSPNet , OCNet , RefineNet 和 DenseASPP等模型虽然分割效果很好，但网络结构笨重，模型参数太多，不利于移植到移动端。

而本文研究了如何在语义分割模型的知识蒸馏中引入结构化信息的损失函数，在不改变模型计算量的情况下，提高小型网络在语义分割任务上的精度。

知识蒸馏

顾名思义，知识蒸馏就是把大网络学习到的知识浓缩到小网络模型中。一般情况下，在相同的数据上训练，模型参数量较大、计算量大的模型往往精度比较高，而用精度高、模型复杂度高的模型即Teacher网络的输出训练Student网络，以期达到使计算量小参数少的小网络精度提升的方法，就是知识蒸馏。

语义分割问题描述

语义分割任务的目标是将图像中的每个像素点划分到对应的类别中。假设目标类别有 ?个 C个，输入图像为 _?_， I，通道数为3（RGB），长宽分别为 ?,?__， H,W，则 ?的 I的尺寸为 _?_× ?__×3。 W×H×3。

将其输入语义分割网络后首先计算得到特征图 ?_,其 F,其尺寸记为 _?′ × ?__′ × ?__， W^′×H′×N， ? N为特征图的通道数，然后使用分类器计算特征图 _? _得到语义图?, 尺寸为?′×?′×?Q, 尺寸为W^′×H′×C再将其长宽上采样到 ?__× ? W×H即可。

结构化知识蒸馏整体架构

本文运用知识蒸馏策略，将大型网络 _T _T（Teacher）的知识转移到小型网络** S** S（Student）中，除了最简单的像素级蒸馏方案外，本文还提出了两种结构化的知识蒸馏方案——成对蒸馏和整体蒸馏，以便将结构化知识从笨重的大型网络传输到小型网络中。

方法

像素蒸馏

由于语义分割本质上是对像素点进行分类，因此可以直接使用知识蒸馏调整小型网络生成的每个像素类别的概率。

本文将大型网络生成的类别概率作为目标训练小型网络，损失函数如下：

其中， _???代表 _小型网络 _? _第 _?个 i个像素点的类别概率， _??_?_代表大 型网络 _? _第 _?_个 i个像素点的类别概率。函数KL(∙)是两个像素点类别概率的相对熵（就是两者的交叉熵相减）。

成对蒸馏

在成对马尔可夫随机场（the pair-wise Markov random field）的启发下，本文提出了一种基于成对像素之间相似性的知识蒸馏策略。（与上一种方法不同，这里的像素指的是特征图上的像素点）

其中，a_ij^s 是小型网络 _?_第 _?_个 像素点 S第i个像素点和 第 _?_个 像素点 之间 第j个像素点之间的相似度值， a_ij^t 是大 型网络 _?_第 _?_个 像素点 和第j个像素点之间的相似度。

由于此处的像素点是特征图 _?_上的 F上的像素点，所以实际计算使用的是特征图上的特征向量，即

其中， ?1 _,?2 即为 f_1,f_2 即为两个像素点对应的特征向量， _||__∙ _|__|_2代表 向量的第二范数。

整体蒸馏

为了学习大型模型的整体特征，本文使用条件生成对抗网络（Conditional Generative Adversarial Networks，简称CGAN）。本文将小型网络作为基于输入图像 ? I的生成器，将其生成的预测图 ?_?作为 负样本，将大型网络的预测图 _??作为 _正样本，通过训练CGAN使 _?? （Q^S）接近 ?? （Q^T） ， 得到的损失函数如下：

公式的本质是Wasserstein distance，它用于计算从 _?? _的概率分布转换到 _?? _的概率分布所需的最短距离，是GAN常用的损失函数，** ?** 是CGAN的判别器。

整合损失函数

将三种知识蒸馏策略得到的损失函数与小型网络自身的损失函数相整合。

其中， ?__?? （l_mc ）即为小型网络本身的损失函数，使用的交叉熵，另外 _?_1 ， _?_2 分别设10和0.1（此处的10、0.1为实验得到）。

实验

本文使用PSPNet 作为大型模型 _? _, 其骨架模型使用ResNet101；使用ResNet18作为本文的小型模型 ?__。

本文使用SGD作为梯度下降算法，设置动量（momentum）为0.9，学习率初始化为0.01，下降率为。

本文使用的数据集：Cityscapes、Camvid和ADE20K。

PI = pixel-wise distillation,像素蒸馏 PA = pair-wise distillation,成对蒸馏 HO =holistic distillation, 整体蒸馏 ImN = initial from the pretrain weight on the ImageNet，在ImageNet上进行与训练过。

W/O=without