【GiantPandaCV引言】 知識(shí)回顧（KR）發(fā)現(xiàn)學(xué)生網(wǎng)絡(luò)深層可以通過利用教師網(wǎng)絡(luò)淺層特征進(jìn)行學(xué)習(xí)，基于此提出了回顧機(jī)制，包括ABF和HCL兩個(gè)模塊，可以在很多分類任務(wù)上得到一致性的提升。
摘要
知識(shí)蒸餾通過將知識(shí)從教師網(wǎng)絡(luò)傳遞到學(xué)生網(wǎng)絡(luò)，但是之前的方法主要關(guān)注提出特征變換和實(shí)施相同層的特征。
知識(shí)回顧Knowledge Review選擇研究教師與學(xué)生網(wǎng)絡(luò)之間不同層之間的路徑鏈接。
簡(jiǎn)單來說就是研究教師網(wǎng)絡(luò)向?qū)W生網(wǎng)絡(luò)傳遞知識(shí)的鏈接方式。
代碼在：https://github.com/Jia-Research-Lab/ReviewKD
KD簡(jiǎn)單回顧
KD最初的蒸餾對(duì)象是logits層，也即最經(jīng)典的Hinton的那篇Knowledge Distillation，讓學(xué)生網(wǎng)絡(luò)和教師網(wǎng)絡(luò)的logits KL散度盡可能小。
隨后FitNets出現(xiàn)開始蒸餾中間層，一般通過使用MSE Loss讓學(xué)生網(wǎng)絡(luò)和教師網(wǎng)絡(luò)特征圖盡可能接近。
Attention Transfer進(jìn)一步發(fā)展了FitNets，提出使用注意力圖來作為引導(dǎo)知識(shí)的傳遞。
PKT(Probabilistic knowledge transfer for deep representation learning)將知識(shí)作為概率分布進(jìn)行建模。
Contrastive representation Distillation(CRD)引入對(duì)比學(xué)習(xí)來進(jìn)行知識(shí)遷移。
以上方法主要關(guān)注于知識(shí)遷移的形式以及選擇不同的loss function，但KR關(guān)注于如何選擇教師網(wǎng)絡(luò)和學(xué)生網(wǎng)絡(luò)的鏈接，一下圖為例：
（a-c）都是傳統(tǒng)的知識(shí)蒸餾方法，通常都是相同層的信息進(jìn)行引導(dǎo)，（d）代表KR的蒸餾方式，可以使用教師網(wǎng)絡(luò)淺層特征來作為學(xué)生網(wǎng)絡(luò)深層特征的監(jiān)督，并發(fā)現(xiàn)學(xué)生網(wǎng)絡(luò)深層特征可以從教師網(wǎng)絡(luò)的淺層學(xué)習(xí)到知識(shí)。
教師網(wǎng)絡(luò)淺層到深層分別對(duì)應(yīng)的知識(shí)抽象程度不斷提高，學(xué)習(xí)難度也進(jìn)行了提升，所以學(xué)生網(wǎng)絡(luò)如果能在初期學(xué)習(xí)到教師網(wǎng)絡(luò)淺層的知識(shí)會(huì)對(duì)整體有幫助。
KR認(rèn)為淺層的知識(shí)可以作為舊知識(shí)，并進(jìn)行不斷回顧，溫故知新。如何從教師網(wǎng)絡(luò)中提取多尺度信息是本文待解決的關(guān)鍵：
提出了Attention based fusion(ABF) 進(jìn)行特征fusion
提出了Hierarchical context loss(HCL) 增強(qiáng)模型的學(xué)習(xí)能力。
Knowledge Review
形式化描述
X是輸入圖像，S代表學(xué)生網(wǎng)絡(luò)，其中代表學(xué)生網(wǎng)絡(luò)各個(gè)層的組成。
Ys代表X經(jīng)過整個(gè)網(wǎng)絡(luò)以后的輸出。代表各個(gè)層中間層輸出。
那么單層知識(shí)蒸餾可以表示為：
M代表一個(gè)轉(zhuǎn)換，從而讓Fs和Ft的特征圖相匹配。D代表衡量?jī)烧叻植嫉木嚯x函數(shù)。
同理多層知識(shí)蒸餾表示為：
以上公式是學(xué)生和教師網(wǎng)絡(luò)層層對(duì)應(yīng)，那么單層KR表示方式為：
與之前不同的是，這里計(jì)算的是從j=1 to i 代表第i層學(xué)生網(wǎng)絡(luò)的學(xué)習(xí)需要用到從第1到i層所有知識(shí)。
同理，多層的KR表示為：
Fusion方式設(shè)計(jì)
已經(jīng)確定了KR的形式，即學(xué)生每一層回顧教師網(wǎng)絡(luò)的所有靠前的層，那么最簡(jiǎn)單的方法是：
直接縮放學(xué)生網(wǎng)絡(luò)最后一層feature，讓其形狀和教師網(wǎng)絡(luò)進(jìn)行匹配，這樣可以簡(jiǎn)單使用一個(gè)卷積層配合插值層完成形狀的匹配過程。這種方式是讓學(xué)生網(wǎng)絡(luò)更接近教師網(wǎng)絡(luò)。
這張圖表示擴(kuò)展了學(xué)生網(wǎng)絡(luò)所有層對(duì)應(yīng)的處理方式，也即按照第一張圖的處理方式進(jìn)行形狀匹配。
這種處理方式可能并不是最優(yōu)的，因?yàn)闀?huì)導(dǎo)致stage之間出現(xiàn)巨大的差異性，同時(shí)處理過程也非常復(fù)雜，帶來了額外的計(jì)算代價(jià)。
為了讓整個(gè)過程更加可行，提出了Attention based fusion , 這樣整體蒸餾變?yōu)椋?/p>
如果引入了fusion的模塊，那整體流程就變?yōu)橄聢D所示：
但是為了更高的效率，再對(duì)其進(jìn)行改進(jìn)：
可以發(fā)現(xiàn)，這個(gè)過程將fusion的中間結(jié)果進(jìn)行了利用，即, 這樣循環(huán)從后往前進(jìn)行迭代，就可以得到最終的loss。
具體來說，ABF的設(shè)計(jì)如下（a）所示，采用了注意力機(jī)制融合特征，具體來說中間的1x1 conv對(duì)兩個(gè)level的feature提取綜合空間注意力特征圖，然后再進(jìn)行特征重標(biāo)定，可以看做SKNet的空間注意力版本。
而HCL Hierarchical context loss 這里對(duì)分別來自于學(xué)生網(wǎng)絡(luò)和教師網(wǎng)絡(luò)的特征進(jìn)行了空間池化金字塔的處理，L2 距離用于衡量?jī)烧咧g的距離。
KR認(rèn)為這種方式可以捕獲不同level的語義信息，可以在不同的抽象等級(jí)提取信息。
實(shí)驗(yàn)
實(shí)驗(yàn)部分主要關(guān)注消融實(shí)驗(yàn)：
第一個(gè)是使用不同stage的結(jié)果：
藍(lán)色的值代表比baseline 69.1更好，紅色代表要比baseline更差。通過上述結(jié)果可以發(fā)現(xiàn)使用教師網(wǎng)絡(luò)淺層知識(shí)來監(jiān)督學(xué)生網(wǎng)絡(luò)深層知識(shí)是有效的。
第二個(gè)是各個(gè)模塊的作用：
源碼
主要關(guān)注ABF， HCL的實(shí)現(xiàn)：
ABF實(shí)現(xiàn)：
class ABF(nn.Module):
def __init__(self, in_channel, mid_channel, out_channel, fuse):
super(ABF, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channel, mid_channel, kernel_size=1, bias=False),
nn.BatchNorm2d(mid_channel),
)
self.conv2 = nn.Sequential(
nn.Conv2d(mid_channel, out_channel,kernel_size=3,stride=1,padding=1,bias=False),
nn.BatchNorm2d(out_channel),
)
if fuse:
self.att_conv = nn.Sequential(
nn.Conv2d(mid_channel*2, 2, kernel_size=1),
nn.Sigmoid(),
)
else:
self.att_conv = None
nn.init.kaiming_uniform_(self.conv1[0].weight, a=1) # pyre-ignore
nn.init.kaiming_uniform_(self.conv2[0].weight, a=1) # pyre-ignore
def forward(self, x, y=None, shape=None, out_shape=None):
n,_,h,w = x.shape
# transform student features
x = self.conv1(x)
if self.att_conv is not None:
# upsample residual features
y = F.interpolate(y, (shape,shape), mode="nearest")
# fusion
z = torch.cat([x, y], dim=1)
z = self.att_conv(z)
x = (x * z[:,0].view(n,1,h,w) + y * z[:,1].view(n,1,h,w))
# output
if x.shape[-1] != out_shape:
x = F.interpolate(x, (out_shape, out_shape), mode="nearest")
y = self.conv2(x)
return y, x
HCL實(shí)現(xiàn)：
def hcl(fstudent, fteacher):
# 兩個(gè)都是list，存各個(gè)stage對(duì)象
loss_all = 0.0
for fs, ft in zip(fstudent, fteacher):
n,c,h,w = fs.shape
loss = F.mse_loss(fs, ft, reduction='mean')
cnt = 1.0
tot = 1.0
for l in [4,2,1]:
if l >=h:
continue
tmpfs = F.adaptive_avg_pool2d(fs, (l,l))
tmpft = F.adaptive_avg_pool2d(ft, (l,l))
cnt /= 2.0
loss += F.mse_loss(tmpfs, tmpft, reduction='mean') * cnt
tot += cnt
loss = loss / tot
loss_all = loss_all + loss
return loss_all
ReviewKD實(shí)現(xiàn)：
class ReviewKD(nn.Module):
def __init__(
self, student, in_channels, out_channels, shapes, out_shapes,
):
super(ReviewKD, self).__init__()
self.student = student
self.shapes = shapes
self.out_shapes = shapes if out_shapes is None else out_shapes
abfs = nn.ModuleList()
mid_channel = min(512, in_channels[-1])
for idx, in_channel in enumerate(in_channels):
abfs.append(ABF(in_channel, mid_channel, out_channels[idx], idx < len(in_channels)-1))
self.abfs = abfs[::-1]
self.to('cuda')
def forward(self, x):
student_features = self.student(x,is_feat=True)
logit = student_features[1]
x = student_features[0][::-1]
results = []
out_features, res_features = self.abfs[0](x[0], out_shape=self.out_shapes[0])
results.append(out_features)
for features, abf, shape, out_shape in zip(x[1:], self.abfs[1:], self.shapes[1:], self.out_shapes[1:]):
out_features, res_features = abf(features, res_features, shape, out_shape)
results.insert(0, out_features)
return results, logit
參考
https://zhuanlan.zhihu.com/p/363994781
