知识蒸馏，在 `Distilling the Knowledge in a Neural Network <https://arxiv.org/abs/1503.02531>`__ 中，压缩模型被训练成模拟预训练的大模型。 这种训练设置也称为"师生（teacher-student）"方式，其中大模型是教师，小模型是学生。
在 `Distilling the Knowledge in a Neural Network <https://arxiv.org/abs/1503.02531>`__\ 中提出了知识蒸馏（KD）的概念, 压缩后的模型被训练去模仿预训练的、较大的模型。 这种训练设置也称为"师生（teacher-student）"方式，其中大模型是教师，小模型是学生。 KD 通常用于微调剪枝后的模型。
.. image:: ../../img/distill.png
.. image:: ../../img/distill.png
:target: ../../img/distill.png
:target: ../../img/distill.png
:alt:
:alt:
用法
用法
^^^^^
^^^^^
...
@@ -19,24 +18,29 @@ PyTorch 代码
...
@@ -19,24 +18,29 @@ PyTorch 代码
.. code-block:: python
.. code-block:: python
from knowledge_distill.knowledge_distill import KnowledgeDistill
for batch_idx, (data, target) in enumerate(train_loader):
kd = KnowledgeDistill(kd_teacher_model, kd_T=5)
data, target = data.to(device), target.to(device)
alpha = 1
optimizer.zero_grad()
beta = 0.8
y_s = model_s(data)
for batch_idx, (data, target) in enumerate(train_loader):
y_t = model_t(data)
data, target = data.to(device), target.to(device)
loss_cri = F.cross_entropy(y_s, target)
optimizer.zero_grad()
output = model(data)
# kd 损失值
loss = F.cross_entropy(output, target)
p_s = F.log_softmax(y_s/kd_T, dim=1)
# 只需要添加以下行来使用知识蒸馏微调模型
p_t = F.softmax(y_t/kd_T, dim=1)
loss = alpha * loss + beta * kd.loss(data=data, student_out=output)
