### 实现知识蒸馏 在PyTorch中实现知识蒸馏可以通过定义教师模型和学生模型来完成。通过让较小的学生模型模仿较大更复杂的教师模型的行为，可以有效地传递知识并提升学生的性能。 #### 定义损失函数 为了使学生模型能够学习到教师模型的知识，通常会引入软目标的概念。这涉及到调整交叉熵损失函数中的温度参数\( T \)，使得概率分布更加平滑[^1]： \[ L_{KD} = (1-\alpha) \cdot CE(\text{logits}_{\text{student}}, y) + \alpha \cdot KL(\frac{\exp(\text{logits}_{\text{teacher}}/T)}{\sum_j \exp(\text{logits}_{j,\text{teacher}}/T)}, \frac{\exp(\text{logits}_{\text{student}}/T)}{\sum_j \exp(\text{logits}_{j,\text{student}}/T)}) \] 其中 \(CE\) 表示标准的交叉熵损失；而 \(KL\) 则用于衡量两个分布之间的差异程度。 ```python import torch.nn.functional as F def kd_loss_fn(logits_student, logits_teacher, labels, alpha=0.5, temperature=4): loss_ce = F.cross_entropy(logits_student, labels) softmax_teachers = F.softmax(logits_teacher / temperature, dim=-1) log_softmax_students = F.log_softmax(logits_student / temperature, dim=-1) loss_kd = F.kl_div(log_softmax_students, softmax_teachers.detach(), reduction='batchmean') total_loss = (1 - alpha) * loss_ce + alpha * temperature ** 2 * loss_kd return total_loss ``` 此代码片段实现了上述提到的知识蒸馏损失函数，并允许设置不同的权重因子 `alpha` 和温度系数 `temperature`. #### 构建训练循环 接下来展示如何在一个简单的训练过程中应用这个自定义的 KD 损失函数: ```python for epoch in range(num_epochs): model.train() running_loss = 0.0 for inputs, targets in train_loader: optimizer.zero_grad() outputs_student = student_model(inputs.to(device)) with torch.no_grad(): outputs_teacher = teacher_model(inputs.to(device)) loss = kd_loss_fn(outputs_student, outputs_teacher, targets.to(device), alpha=args.alpha, temperature=args.temperature) loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) ``` 这段代码展示了完整的训练过程，包括前向传播、计算损失以及反向传播更新参数等操作。