一文弄懂模型知识蒸馏、剪枝

### 深度学习中模型的知识蒸馏与剪枝技术原理及应用

#### 知识蒸馏的技术原理及应用

知识蒸馏是一种有效的模型压缩策略，旨在将复杂的教师模型（teacher model）所蕴含的知识迁移到较为简单的学生模型（student model）。这种方法不仅能够保持较高的预测准确性，还能显著减小模型尺寸并提升运行效率。具体来说，在训练过程中，学生模型不仅要拟合原始数据集上的标签信息，还要尽可能模仿教师模型给出的概率分布，即所谓的“暗知识”或软标签[^1]。

为了实现这一点，通常会采用温度缩放机制调整softmax函数的输出，使得教师网络产生的概率分布更加平滑，便于学生更好地捕捉其特征表示能力。此外，还可以引入额外损失项来强化这种迁移过程的效果，比如基于中间层激活值的一致性约束等[^5]。

import torch.nn.functional as F

def knowledge_distillation_loss(student_logits, teacher_logits, temperature=2.0):
    soft_student = F.softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)

    loss_kd = F.kl_div(
        input=F.log_softmax(student_logits / temperature, dim=-1),
        target=F.softmax(teacher_logits / temperature, dim=-1),
        reduction='batchmean'
    ) * (temperature ** 2)

    return loss_kd

#### 剪枝的方法论及其应用场景

相比之下，剪枝则是另一种不同的模型简化手段，主要关注于移除那些对整体性能贡献较小甚至可以忽略不计的部分——通常是连接权值接近零的位置。通过对神经元间联系强度进行评估筛选，并逐步去除冗余组件，最终得到一个更为紧凑高效的版本[^3]。

实际操作时，一般先完成一次完整的预训练阶段；接着依据设定的标准挑选出待修剪的目标节点/边；最后重新微调剩余结构直至满足预期指标为止。值得注意的是，尽管此法能在一定程度上缓解过拟合现象的发生几率，但也可能导致泛化能力下降等问题出现，因此需谨慎对待参数设置环节[^2]。

from functools import partial
import numpy as np

def prune_weights(model, pruning_ratio=0.2):
    all_params = []
    
    for name, param in model.named_parameters():
        if 'weight' in name and not ('bn' in name or 'bias' in name):  
            all_params.append((name, param.data.cpu().numpy()))
            
    flat_params = np.concatenate([p.flatten() for _, p in all_params])
    threshold = np.percentile(abs(flat_params), q=(pruning_ratio*100))
    
    with torch.no_grad():
        for layer_name, weights in all_params:
            mask = abs(weights) >= threshold
            pruned_tensor = torch.from_numpy(mask.astype(int)).cuda()
            getattr(model, '.'.join(layer_name.split('.')[:-1]))._parameters[layer_name.split('.')[-1]].mul_(pruned_tensor)