Overhaul Distillation(ICCV 2019)原理与代码解析

paper:A Comprehensive Overhaul of Feature Distillation

official implementation:GitHub - clovaai/overhaul-distillation: Official PyTorch implementation of "A Comprehensive Overhaul of Feature Distillation" (ICCV 2019)

本文的创新点

本文研究了知识蒸馏的各个方面,并提出了一种新的特征蒸馏方法,使蒸馏损失在教师特征变换、学生特征变换、特征蒸馏位置、距离函数各方面之间协同作用。具体来说,本文提出的蒸馏损失包括一个新设计的margin relu特征变换方法、一个新的蒸馏位置、以及一个partial L2距离函数。在ImageNet中,本文提出的方法使得ResNet-50取得了21.65%的top-1 error,优于教师网络ResNet-152的精度。

方法介绍 

蒸馏位置

激活函数是神经网络的重要组成部分,它使网络具有了非线性。但之前的大部分蒸馏方法都没有考虑到激活函数,蒸馏位置大都在于某各个layer或某个block的尾端,却没有考虑和激活函数如ReLU的关系。本文提出的方法中,蒸馏位置位于某个layer的尾端和第一个ReLU之间,如下图所示

Overhaul Distillation(ICCV 2019)原理与代码解析_第1张图片 

pre-ReLU的位置可以使学生接触到教师模型通过ReLU之间的信息,避免了信息的分解和丢失。

损失函数

由于蒸馏位置是在ReLU之前,因此特征中的正值包含教师会利用的信息而负值没有,如果教师网络中的值是正值,学生网络应该生成和教师一样的值,如果教师网络中是负值,学生也应该生成负值从而使得激活状态和教师一致。因此作者提出的教师变换函数,保存正值同时有一个负值margin 

其中 \(m\) 是一个小于0的margin值,作者取名为margin ReLu。\(m\) 的具体值定义为每个通道的负响应值的期望,如下

 

\(m\) 一方面可以在训练过程中直接计算,如https://github.com/clovaai/overhaul-distillation/issues/7所示。也可以通过前一个BN层的参数来计算,作者在附录中给出了具体计算方法。

对于一个通道 \(\mathcal{C}\) 和教师特征 \(F^{i}_{t}\) 的第 \(i\) 个元素,该通道的margin值 \(m_{c}\) 为训练图片的期望值即式(3)。通常我们不知道 \(F^{i}_{t}\) 的分布,所以只能通过训练过程中的平均值来得到期望。但是ReLU前的BN层决定了一个batch中的特征 \(F^{i}_{t}\) 的分布,BN层将每个通道的特征归一化为均值 \(\mu\) 方差 \(\sigma\) 的高斯分布,即

每个通道的均值方差 \((\mu,\sigma)\) 对应BN层的参数 \((\beta,\gamma)\),因此利用 \(F^{i}_{t}\) 的分布可以直接计算边际值

 

利用高斯分布的概率密度函数pdf进行积分就可以得到期望,其中范围小于0。积分的结果可以通过正太分布的cdf累积分布函数 \(\Phi(\cdot)\) 进行简单的表示。

 

在官方实现中,也是通过这种方式即式(10)来计算margin值的。

因为蒸馏的位置是在ReLU函数前,negative response没有经过ReLU的过滤,因此蒸馏损失函数需要考虑到ReLU。在教师特征中,positive response实际上被网络使用,这意味着教师的正响应应该通过具体的值来传递,但负响应却不是。对于教师的负响应,如果学生的响应值高于目标值应该降低,而如果低于目标值不需要增加,因为不管具体值是多少都会被ReLU过滤掉。因此,本文提出了partial L2 distance函数,如下

完整的蒸馏损失函数如下

 

其中 \(\sigma_{m_{c}}\) 是教师转换函数margin ReLU,\(r\) 是是学生转换函数1x1 conv + BN,\(d_{p}\) 是距离函数patial L2 distance。

实验结果

在CIFAR-100数据集上,不同的教师网络和学生网络的结果如表2所示

Overhaul Distillation(ICCV 2019)原理与代码解析_第2张图片

不同的教师-学生网络组合,本文的方法和其它蒸馏方法的结果对比如下,可以看出,在所有组合下,本文提出的方法都得到了最低的error。

Overhaul Distillation(ICCV 2019)原理与代码解析_第3张图片

在ImageNet数据集上和其它方法的对比如表4,可以看出本文的方法error也是最低的。 

Overhaul Distillation(ICCV 2019)原理与代码解析_第4张图片

代码解析

实现代码主要在distiller.py中,本文的第一个创新点在蒸馏的位置,即ReLU前,实现如下 

t_feats, t_out = self.t_net.extract_feature(x, preReLU=True)
s_feats, s_out = self.s_net.extract_feature(x, preReLU=True)

学生特征的转换为1x1卷积+BN,即实现中的self.Connectors,具体实现如下

def build_feature_connector(t_channel, s_channel):
    C = [nn.Conv2d(s_channel, t_channel, kernel_size=1, stride=1, padding=0, bias=False),
         nn.BatchNorm2d(t_channel)]

    for m in C:
        if isinstance(m, nn.Conv2d):
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()

    return nn.Sequential(*C)

教师特征的转换为本文提出的margin ReLU,边际margin值的计算如下,即上述的式(10)

def get_margin_from_BN(bn):
    margin = []
    std = bn.weight.data
    mean = bn.bias.data
    for (s, m) in zip(std, mean):
        s = abs(s.item())
        m = m.item()
        if norm.cdf(-m / s) > 0.001:
            margin.append(- s * math.exp(- (m / s) ** 2 / 2) / math.sqrt(2 * math.pi) / norm.cdf(-m / s) + m)
        else:
            margin.append(-3 * s)

    return torch.FloatTensor(margin).to(std.device)

蒸馏损失函数实现如下,其中第一行就是教师特征的转换函数,即式(2)

def distillation_loss(source, target, margin):
    target = torch.max(target, margin)
    loss = torch.nn.functional.mse_loss(source, target, reduction="none")
    loss = loss * ((source > target) | (target > 0)).float()
    return loss.sum()

你可能感兴趣的:(知识蒸馏-分类,python,算法,人工智能,深度学习,知识蒸馏)