知识蒸馏

最新推荐文章于 2024-01-21 03:37:34 发布

影醉阏轩窗

最新推荐文章于 2024-01-21 03:37:34 发布

阅读量372

点赞数

文章标签：网络深度学习 python 机器学习人工智能

本文链接：https://blog.csdn.net/u011046017/article/details/115521535

版权

知识蒸馏

一. Distilling the Knowledge in a Neural Network

知识蒸馏的开端之作，简单叙述蒸馏过程：

先训练一个大网络，比如Resnet50用于分类任务
搭建一个小网络训练结构，比如mobilenetV2
训练小网络的同时推理大网络，大网络的结果去指导小网络(KDLoss用于估计分布的相似性)

类似的代码：链接地址

类似的文章：链接地址

比较简单的过程：

# 教师输出和学生输出得到loss1，学生输出和label得到loss2，按一定比例结合进行反向传播
def loss_fn_kd(outputs, labels, teacher_outputs, params):
    """
    Compute the knowledge-distillation (KD) loss given outputs, labels.
    "Hyperparameters": temperature and alpha
    NOTE: the KL Divergence for PyTorch comparing the softmaxs of teacher
    and student expects the input tensor to be log probabilities! See Issue #2
    """
    alpha = params.alpha
    T = params.temperature
    KD_loss = nn.KLDivLoss()(F.log_softmax(outputs/T, dim=1),
                             F.softmax(teacher_outputs/T, dim=1)) * (alpha * T * T) + \
              F.cross_entropy(outputs, labels) * (1. - alpha)

    return KD_loss

二. Fast Human Pose Estimation Pytorch

论文：链接地址

代码：链接地址

论文没有实质的创新，KDLoss直接使用MSE对Heatmap进行分布相似估计，正常Loss也使用MSE，按一定比例结核即可

注释：看见有人说蒸馏必须网络结构类似，不然效果反而会下降(待尝试)

# 关键点不可见的情况下只进行KDLoss，可见的情况下进行KDLoss和正常训练Loss
for j in range(0, len(output)):
  	_output = output[j]
  	for i in range(gtmask.shape[0]):
          if gtmask[i] < 0.1:
          # unlabeled data, gtmask=0.0, kdloss only
          # need to dividen train_batch to keep number equal
          kdloss_unlabeled += criterion(_output[i,:,:,:], toutput[i, :,:,:])/train_batch
      	else:
          # labeled data: kdloss + gtloss
          gtloss += criterion(_output[i,:,:,:], target_var[i, :,:,:])/train_batch
          kdloss += criterion(_output[i,:,:,:], toutput[i,:,:,:])/train_batch

loss_labeled = kdloss_alpha * (kdloss) + (1 - kdloss_alpha)*gtloss
total_loss   = loss_labeled + unkdloss_alpha * kdloss_unlabeled