知识蒸馏

term_id: knowledge_distillation

Category: training_techniques

Definition

知识蒸馏是一种机器学习方法，用于将庞大复杂的神经网络（教师模型）压缩为更小、更高效的网络（学生模型）。学生模型经过训练以模仿教师模型的行为。

Summary

知识蒸馏是一种模型压缩技术，其中较小的学生模型通过学习模仿较大的教师模型的行为来工作。

Key Concepts

  • 师生模型
  • 模型压缩
  • 软标签
  • 效率

Use Cases

  • 在边缘设备上部署模型
  • 降低推理延迟
  • 降低云计算成本

Code Example

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

import torch
import torch.nn as nn

def distillation_loss(student_logits, teacher_logits, temperature=2.0):
    T = temperature
    student_probs = nn.functional.softmax(student_logits / T, dim=1)
    teacher_probs = nn.functional.softmax(teacher_logits / T, dim=1)
    return nn.functional.kl_div(
        nn.functional.log_softmax(student_logits / T, dim=1),
        teacher_probs,
        reduction='batchmean'
    ) * (T * T)

Related