模型压缩技术详解：剪枝、量化与知识蒸馏

模型压缩技术详解：剪枝、量化与知识蒸馏

引用

CSDN

1.

https://blog.csdn.net/u010618499/article/details/137479217

模型压缩是当前AI领域的重要研究方向，其主要目的是在保证模型精度的同时，减少模型的计算量和存储需求。本文将详细介绍三种常见的模型压缩方法：剪枝、量化和知识蒸馏，并通过具体的代码示例帮助读者理解这些技术的实现细节。

1. 剪枝（Pruning）

剪枝是一种通过减少模型中的参数量来进行模型压缩的技术，可以在保证一定精度的情况下减少模型的内存占用和硬件消耗。

PyTorch剪枝实现

PyTorch提供了多种剪枝规则，包括RandomUnstructured、L1Unstructured、RandomStructured、LnStructured和CustomFromMask。下面是一个使用L1Unstructured方法进行全局剪枝的示例：

# 实例化神经网络
model = LeNet()
# 配置可剪枝的网络层和参数名
parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
    (model.fc2, 'weight'),
    (model.fc3, 'weight'),
)
# 全局剪枝，采用L1Unstructured的方法，剪去0.2的参数量
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2,
)
# 查看剪枝后的模型参数
print(model.state_dict())

2. 量化（Quantization）

模型量化是一种压缩网络参数的方式，它将神经网络的参数（weight）、特征图（activation）等原本用浮点表示的量值换用定点（整型）表示，同时尽可能减少计算精度损失。

PyTorch量化实现

PyTorch提供了三种量化方式：量化感知训练（QAT）、训练后动态量化和训练后校正量化。

"""训练后动态量化"""
import torch.quantization
quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

"""训练后校正量化"""
# 设置服务器(x86)的量化配置
deploymentmyModel.qconfig = torch.quantization.get_default_config('fbgemm')
torch.quantization.prepare(myModel, inplace=True)
torch.quantization.convert(myModel, inplace=True)

"""量化感知训练"""
# 设置量化感知训练服务器(x86)的量化配置
qat_model.qconfig=torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(qat_model, inplace=True)
# 转换为量化版本，删除dropout
epochquantized_model=torch.quantization.convert(qat_model.eval(), inplace=False)

量化方法分类

根据量化方案分类：

• 在线量化（QAT）：在量化的同时结合反向传播对模型权重进行调整
• 离线量化（PTQ）：在已训练的模型上进行量化校准

根据量化公式分类：

• 线性量化：浮点数与定点数之间的转换公式为 Q = R / S + Z
• 对称量化：量化前后的0点是对齐的，适合对分布良好且均值为0的参数进行量化
• 非对称量化：量化前后0点不对齐，需要额外记录一个offset，适合对activation进行量化

Block-wise量化

为了避免异常值的影响，可以将输入tensor分割成一个个block，每个block单独进行量化，有单独的scale和zero，从而减少量化的精度损失。

3. 知识蒸馏（Knowledge Distillation）

知识蒸馏是一种技术，可以将知识从计算成本较高的大型模型转移到较小的模型，而不会失去有效性。具体的方法是在训练小模型时，在损失函数中添加额外的损失函数。

损失函数添加方式

1. 输出层的差异损失

这种方法是对比学生网络和教师网络在输出层上的概率分布之间的差异。

optimizer.zero_grad()
ce_loss = nn.CrossEntropyLoss()
# 冻结教师网络的权重，计算教师网络的输出层结果
with torch.no_grad(): 
    teacher_logits = teacher(inputs) 
# 学生网络的输出层结果 
student_logits = student(inputs) 
# 计算软目标损失（soft targets loss），考虑温度系数T的影响，计算学生网络推理结果和教师网络在推理结果上的差异
soft_targets = nn.functional.softmax(teacher_logits / T, dim=-1)
soft_prob = nn.functional.log_softmax(student_logits / T, dim=-1)
soft_targets_loss = torch.sum(soft_targets * (soft_targets.log() - soft_prob)) / soft_prob.size()[0] * (T**2) 
# 学生网络的交叉熵损失
label_loss = ce_loss(student_logits, labels) 
# 将软目标损失添加到交叉熵损失中
loss = soft_target_loss_weight * soft_targets_loss + ce_loss_weight * label_loss 
loss.backward()
optimizer.step()

2. 隐藏层的相似度损失（余弦相似度损失）

这种方法是对比学生网络和教师网络在输出层上的概率分布之间的差异。

optimizer.zero_grad()
ce_loss = nn.CrossEntropyLoss()
cosine_loss = nn.CosineEmbeddingLoss()
# 冻结教师网络的权重，只记录隐藏层的结果（一般是在分类器之前，flatten之后的输出结果）
with torch.no_grad(): 
    _, teacher_hidden_representation = teacher(inputs)
# 学生网络的推理结果，包含输出层和隐藏层的结果
student_logits, student_hidden_representation = student(inputs) 
# 计算学生网络和教师网络在隐藏层的余弦损失
hidden_rep_loss = cosine_loss(student_hidden_representation, teacher_hidden_representation, target=torch.ones(inputs.size(0)).to(device))
# 学生网络的交叉熵损失
label_loss = ce_loss(student_logits, labels)
# 将余弦相似度损失添加到交叉熵损失中
loss = hidden_rep_loss_weight * hidden_rep_loss + ce_loss_weight * label_loss
loss.backward()
optimizer.step()

3. 中间层的回归损失（均方误差，MSE）

optimizer.zero_grad()
ce_loss = nn.CrossEntropyLoss()
mse_loss = nn.MSELoss()
# 冻结教师网络的权重，只记录中间层的特征（一般是卷积层的输出结果）
with torch.no_grad(): 
    _, teacher_feature_map = teacher(inputs)
# 学生网络的推理结果，包含输出层的结果和中间层的特征
student_logits, regressor_feature_map = student(inputs) 
# 计算学生网络和教师网络在中间层的均方误差
hidden_rep_loss = mse_loss(regressor_feature_map, teacher_feature_map)
# 学生网络的交叉熵损失
label_loss = ce_loss(student_logits, labels)
# 将中间层的均方误差添加到交叉熵损失中
loss = feature_map_weight * hidden_rep_loss + ce_loss_weight * label_loss
loss.backward()
optimizer.step()