知识蒸馏的蒸馏丧失方法代码总结(包罗：基于logits的方法：KLDiv，dist，dk - ToB企服应用市场:ToB评测及商务社交产业平台

1.利用教师模型的输出概率（基于logits的方法)
2.利用教师模型的中间特征（基于提示的方法）
3.自蒸馏方法

复制代码

import torch.nn as nn
def cosine_similarity(a, b, eps=1e-8):
return (a * b).sum(1) / (a.norm(dim=1) * b.norm(dim=1) + eps)
def pearson_correlation(a, b, eps=1e-8):
return cosine_similarity(a - a.mean(1).unsqueeze(1),
b - b.mean(1).unsqueeze(1), eps)
def inter_class_relation(soft_student_outputs, soft_teacher_outputs):
return 1 - pearson_correlation(soft_student_outputs, soft_teacher_outputs).mean()
def intra_class_relation(soft_student_outputs, soft_teacher_outputs):
return inter_class_relation(soft_student_outputs.transpose(0, 1), soft_teacher_outputs.transpose(0, 1))
class DIST(nn.Module):
def __init__(self, beta=1.0, gamma=1.0, temp=1.0):
super(DIST, self).__init__()
self.beta = beta
self.gamma = gamma
self.temp = temp
def forward(self, student_preds, teacher_preds, **kwargs):
soft_student_outputs = (student_preds / self.temp).softmax(dim=1)
soft_teacher_outputs = (teacher_preds / self.temp).softmax(dim=1)
inter_loss = self.temp ** 2 * inter_class_relation(soft_student_outputs, soft_teacher_outputs)
intra_loss = self.temp ** 2 * intra_class_relation(soft_student_outputs, soft_teacher_outputs)
kd_loss = self.beta * inter_loss + self.gamma * intra_loss
return kd_loss

复制代码

import torch.nn as nn
import torch.nn.functional as F
# loss = alpha * hard_loss + (1-alpha) * kd_loss，此处是单单的kd_loss
class KLDiv(nn.Module):
def __init__(self, temp=1.0):
super(KLDiv, self).__init__()
self.temp = temp
def forward(self, student_preds, teacher_preds, **kwargs):
soft_student_outputs = F.log_softmax(student_preds / self.temp, dim=1)
soft_teacher_outputs = F.softmax(teacher_preds / self.temp, dim=1)
kd_loss = F.kl_div(soft_student_outputs, soft_teacher_outputs, reduction="none").sum(1).mean()
kd_loss *= self.temp ** 2
return kd_loss

复制代码

import torch
import torch.nn as nn
import torch.nn.functional as F
def dkd_loss(logits_student, logits_teacher, target, alpha, beta, temperature):
gt_mask = _get_gt_mask(logits_student, target)
other_mask = _get_other_mask(logits_student, target)
pred_student = F.softmax(logits_student / temperature, dim=1)
pred_teacher = F.softmax(logits_teacher / temperature, dim=1)
pred_student = cat_mask(pred_student, gt_mask, other_mask)
pred_teacher = cat_mask(pred_teacher, gt_mask, other_mask)
log_pred_student = torch.log(pred_student)
tckd_loss = (
F.kl_div(log_pred_student, pred_teacher, reduction='batchmean')
* (temperature ** 2)
)
pred_teacher_part2 = F.softmax(
logits_teacher / temperature - 1000.0 * gt_mask, dim=1
)
log_pred_student_part2 = F.log_softmax(
logits_student / temperature - 1000.0 * gt_mask, dim=1
)
nckd_loss = (
F.kl_div(log_pred_student_part2, pred_teacher_part2, reduction='batchmean')
* (temperature ** 2)
)
return alpha * tckd_loss + beta * nckd_loss
def _get_gt_mask(logits, target):
target = target.reshape(-1)
mask = torch.zeros_like(logits).scatter_(1, target.unsqueeze(1), 1).bool()
return mask
def _get_other_mask(logits, target):
target = target.reshape(-1)
mask = torch.ones_like(logits).scatter_(1, target.unsqueeze(1), 0).bool()
return mask
def cat_mask(t, mask1, mask2):
t1 = (t * mask1).sum(dim=1, keepdims=True)
t2 = (t * mask2).sum(1, keepdims=True)
rt = torch.cat([t1, t2], dim=1)
return rt
class DKD(nn.Module):
def __init__(self, alpha=1., beta=2., temperature=1.):
super(DKD, self).__init__()
self.alpha = alpha
self.beta = beta
self.temperature = temperature
def forward(self, z_s, z_t, **kwargs):
target = kwargs['target']
if len(target.shape) == 2: # mixup / smoothing
target = target.max(1)[1]
kd_loss = dkd_loss(z_s, z_t, target, self.alpha, self.beta, self.temperature)
return kd_loss

复制代码

https://github.com/dvlab-research/ReviewKD

复制代码

https://zhuanlan.zhihu.com/p/603748226?utm_id=0

复制代码