paper:A Gift From Knowledge Distillation: Fast Optimization, Network Minimization and Transfer Learning
code:https://github.com/HobbitLong/RepDistiller/blob/master/distiller_zoo/FSP.py
背景
深度神经网络DNN逐层生成特征。更高层的特征更接近于任务的有用特征。如果我们把DNN的输入看作问题,把输出看作答案,我们就可以把DNN中间生成的特征看作是求解过程中的中间结果。根据这一想法,FitNets可以让学生网络简单地模拟教师网络的中间结果。然而在DNN中,有许多方法或途径来解决从输入生成输出的问题。因此,模拟教师网络生成的特征对学生网络来说是一个硬约束hard constraint。就人而言,老师解释问题的解决过程,学生学习解决问题的流程。当输入特定的问题时,学生网络不一定需要学习中间输出,但当遇到特定类型的问题时,学生网络可以学习这一类问题的通用解决方法。因此作者认为,对于知识蒸馏中的教师网络,演示问题的解决过程比演示中间结果具有更好的泛化性。
本文的创新点
本文将神经网络中层与层之间的信息流动定义为需要蒸馏的知识,并通过计算两个特征层之间的内积来得到这种知识。当将这种层之间的流动作为知识传递给学生网络时,作者通过实验得到了三个结论:
从教师网络学习这种蒸馏知识的学生网络比原始网络的优化(收敛)速度快得多。
学习这种蒸馏知识的学生网络比原始网络的性能更好。
即使教师网络是在一个不同的任务或数据集上训练得到的,学生网络也可以从教师网络中学习到这种知识,并且比从头训练的效果更好。
下图是本文提出的知识蒸馏方法的概念图
![](https://i-blog.csdnimg.cn/blog_migrate/9c1dbc4a33cab088bdd2075969833a91.png)
本文的贡献如下:
提出了一种知识蒸馏的新方法。
这种知识对于快速优化非常有用。
利用所提出的蒸馏知识定义网络的初始权重可以提高小模型的性能。
即使学生网络接受了与教师网络不同的训练任务,所提出的蒸馏知识也能提高学生网络的表现。
方法介绍
作者设计了网络中两个相邻层之间的FSP(flow of solution procedure)矩阵来表示问题的求解过程,对于挑选的层1输出的feature map表示为 \(F^{1}\in \mathbb{R}^{h\times w\times m}\),其中 \(h,w,m\) 分别表示特征图的高、宽、通道数。层2表示为 \(F^{2}\in \mathbb{R}^{h\times w\times n}\),则FSP矩阵 \(G\in \mathbb{R}^{m\times n}\) 可通过下式求得
![](https://i-blog.csdnimg.cn/blog_migrate/59d4dd0abab99ee045d9cfbfb7214689.png)
其中 \(x\) 表示输入图片,\(W\) 表示网络权重参数。
对于残差网络,网络在一些位置的spatial size发生变化,我们选择教师网络和学生网络对应位置具有相同spatial size的特征图来生成FSP matrix,下图是一个示例
![](https://i-blog.csdnimg.cn/blog_migrate/c78923e6c78a2dfcad741a33dcf406ca.png)
计算教师网络和学生网络对应FSP矩阵的L2损失,完整是损失函数如下
![](https://i-blog.csdnimg.cn/blog_migrate/4fab809ddb42bee5c02b6f53696f4c09.png)
其中 \(\lambda_{i}\) 表示每一对FSP矩阵损失的权重,文中设定所有层计算的FSP之间的损失权重相等。\(N\) 表示所有的采样点。
代码解析
forward函数的输入g_s和g_t分别表示学生网络和教师网络中所有用来计算FSP矩阵的层,在compute_fsp中每一层都与相邻层计算fsp矩阵,注意这里的相邻并不是说在原始网络中这两层的相邻的。这里相邻层之间计算fsp矩阵需要保证spatial size相等,如果不相等通过自适应平均池化使之相等。
from __future__ import print_function
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
class FSP(nn.Module):
"""A Gift from Knowledge Distillation:
Fast Optimization, Network Minimization and Transfer Learning"""
def __init__(self, s_shapes, t_shapes):
super(FSP, self).__init__()
assert len(s_shapes) == len(t_shapes), 'unequal length of feat list'
s_c = [s[1] for s in s_shapes]
t_c = [t[1] for t in t_shapes]
if np.any(np.asarray(s_c) != np.asarray(t_c)):
raise ValueError('num of channels not equal (error in FSP)')
def forward(self, g_s, g_t):
# [(64,32,32,32),(64,64,32,32),(64,128,16,16),(64,256,8,8)]
# [(64,32,32,32),(64,64,32,32),(64,128,16,16),(64,256,8,8)]
s_fsp = self.compute_fsp(g_s)
t_fsp = self.compute_fsp(g_t)
loss_group = [self.compute_loss(s, t) for s, t in zip(s_fsp, t_fsp)]
return loss_group
@staticmethod
def compute_loss(s, t):
return (s - t).pow(2).mean()
@staticmethod
def compute_fsp(g):
fsp_list = []
for i in range(len(g) - 1):
bot, top = g[i], g[i + 1] # (64,32,32,32),(64,64,32,32)
b_H, t_H = bot.shape[2], top.shape[2]
if b_H > t_H:
bot = F.adaptive_avg_pool2d(bot, (t_H, t_H))
elif b_H < t_H:
top = F.adaptive_avg_pool2d(top, (b_H, b_H))
else:
pass
bot = bot.unsqueeze(1) # (64,1,32,32,32)
top = top.unsqueeze(2) # (64,64,1,32,32)
bot = bot.view(bot.shape[0], bot.shape[1], bot.shape[2], -1) # (64,1,32,1024)
top = top.view(top.shape[0], top.shape[1], top.shape[2], -1) # (64,64,1,1024)
fsp = (bot * top).mean(-1) # (64,64,32,1024)->(64,64,32)
fsp_list.append(fsp)
return fsp_list