深度学习模型压缩：从剪枝到知识蒸馏

深度学习模型压缩从剪枝到知识蒸馏1. 技术分析1.1 模型压缩方法对比方法压缩比精度损失计算开销适用场景剪枝2x-10x1-5%低所有模型量化2x-4x0.5-3%低推理优化知识蒸馏可变可忽略中分类/检测低秩分解2x-5x1-3%中CNN/全连接1.2 压缩效果评估指标定义测量方法压缩比原始/压缩后大小参数数量比加速比原始/压缩后推理时间时间测量精度损失原始-压缩后精度指标对比2. 核心功能实现2.1 模型剪枝import torch import torch.nn as nn import numpy as np class ModelPruner: 模型剪枝器 def __init__(self, model: nn.Module): self.model model self.original_params self._count_params() self.masks {} def _count_params(self) - int: 统计参数量 return sum(p.numel() for p in self.model.parameters()) def magnitude_pruning(self, sparsity: float 0.5): 幅度剪枝 Args: sparsity: 稀疏度0-1之间的比例 for name, param in self.model.named_parameters(): if weight in name: # 计算幅度 magnitudes torch.abs(param.data) # 计算阈值 threshold torch.quantile(magnitudes, sparsity) # 创建掩码 mask magnitudes threshold # 应用掩码 param.data * mask.float() # 保存掩码 self.masks[name] mask def random_pruning(self, sparsity: float 0.5): 随机剪枝 for name, param in self.model.named_parameters(): if weight in name: mask torch.rand_like(param.data) sparsity param.data * mask.float() self.masks[name] mask def structured_pruning(self, methodchannel, amount0.3): 结构化剪枝 Args: method: channel, filter, block amount: 剪枝比例 for name, module in self.model.named_modules(): if isinstance(module, nn.Conv2d): # 计算每个通道的重要性 importance module.weight.data.abs().sum(dim(0, 2, 3)) # 确定要剪枝的通道 num_to_prune int(len(importance) * amount) threshold torch.kthvalue(importance, num_to_prune)[0] # 创建掩码 mask importance threshold # 应用剪枝 module.weight.data * mask.view(1, -1, 1, 1).float() def apply_masks(self): 应用保存的掩码 for name, param in self.model.named_parameters(): if name in self.masks: param.data * self.masks[name].float() def get_sparsity(self) - float: 获取当前稀疏度 total_zeros 0 total_elements 0 for param in self.model.parameters(): total_elements param.numel() total_zeros (param.data 0).sum().item() return total_zeros / total_elements def stats(self) - dict: 获取剪枝统计 current_params self._count_params() return { original_params: self.original_params, current_params: current_params, compression_ratio: self.original_params / current_params, sparsity: self.get_sparsity() } # 使用示例 model torch.nn.Sequential( torch.nn.Linear(784, 256), torch.nn.ReLU(), torch.nn.Linear(256, 10) ) pruner ModelPruner(model) pruner.magnitude_pruning(sparsity0.5) print(f剪枝统计: {pruner.stats()})2.2 模型量化import torch import torch.nn as nn class Quantizer: 模型量化器 staticmethod def dynamic_quantize(model: nn.Module): 动态量化权重int8激活fp32 return torch.quantization.quantize_dynamic( model, {nn.Linear, nn.Conv2d}, dtypetorch.qint8 ) staticmethod def static_quantize(model: nn.Module, calibration_data): 静态量化 model.eval() # Fuse模块 model torch.quantization.fuse_modules( model, [[0, 1]] # Conv BN ReLU ) # 设置量化配置 model.qconfig torch.quantization.get_default_qconfig(fbgemm) torch.quantization.prepare(model, inplaceTrue) # 校准 print(校准中...) model.eval() with torch.no_grad(): for i, (data, _) in enumerate(calibration_data): if i 100: break model(data) # 转换 quantized_model torch.quantization.convert(model, inplaceFalse) return quantized_model staticmethod def post_training_quantize(model: nn.Module, calibration_data): 训练后量化 # 配置 model.eval() model.fuse_model() # PTQ torch.quantization.quantize_dynamic( model, {nn.Linear, nn.Conv2d}, dtypetorch.qint8, inplaceTrue ) return model def compare_quantization(): 对比不同量化方法 model SimpleModel() # 原始模型大小 original_size sum(p.numel() * p.element_size() for p in model.parameters()) # 动态量化 dynamic_model Quantizer.dynamic_quantize(model) dynamic_size sum( p.numel() * p.element_size() for p in dynamic_model.parameters() ) print(f原始模型: {original_size / 1024 / 1024:.2f} MB) print(f动态量化: {dynamic_size / 1024 / 1024:.2f} MB) print(f压缩比: {original_size / dynamic_size:.2f}x)2.3 知识蒸馏import torch import torch.nn as nn import torch.nn.functional as F class Distiller: 知识蒸馏器 def __init__(self, teacher: nn.Module, student: nn.Module, temperature: float 4.0, alpha: float 0.5): self.teacher teacher self.student student self.temperature temperature self.alpha alpha # 冻结教师模型 for param in self.teacher.parameters(): param.requires_grad False def distillation_loss(self, student_logits, teacher_logits, labels): 蒸馏损失 结合硬标签损失和软标签损失 # 软标签损失KL散度 soft_teacher F.softmax(teacher_logits / self.temperature, dim1) soft_student F.log_softmax(student_logits / self.temperature, dim1) soft_loss F.kl_div( soft_student, soft_teacher, reductionbatchmean ) * (self.temperature ** 2) # 硬标签损失 hard_loss F.cross_entropy(student_logits, labels) # 组合损失 return self.alpha * soft_loss (1 - self.alpha) * hard_loss def train_student(self, train_loader, optimizer, epochs10): 训练学生模型 self.teacher.eval() self.student.train() for epoch in range(epochs): total_loss 0 correct 0 total 0 for data, labels in train_loader: optimizer.zero_grad() # 教师输出 with torch.no_grad(): teacher_logits self.teacher(data) # 学生输出 student_logits self.student(data) # 计算损失 loss self.distillation_loss(student_logits, teacher_logits, labels) loss.backward() optimizer.step() total_loss loss.item() _, predicted student_logits.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() accuracy 100. * correct / total avg_loss total_loss / len(train_loader) print(fEpoch {epoch1}/{epochs}, Loss: {avg_loss:.4f}, Acc: {accuracy:.2f}%) def evaluate(self, test_loader): 评估学生模型 self.student.eval() correct 0 total 0 with torch.no_grad(): for data, labels in test_loader: outputs self.student(data) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() return 100. * correct / total # 使用示例 teacher create_large_model() # 教师模型 student create_small_model() # 学生模型 distiller Distiller(teacher, student, temperature4.0, alpha0.5) optimizer torch.optim.Adam(student.parameters(), lr0.001) distiller.train_student(train_loader, optimizer, epochs20) accuracy distiller.evaluate(test_loader) print(f学生模型测试准确率: {accuracy:.2f}%)2.4 低秩分解import torch import torch.nn as nn from scipy.linalg import svd class LowRankDecomposition: 低秩分解 staticmethod def decompose_linear(layer: nn.Linear, rank_ratio: float 0.5): 分解Linear层 将 W 分解为 W U V 其中 U 的形状是 (out_features, rank) V 的形状是 (rank, in_features) W layer.weight.data.numpy() original_shape W.shape # SVD分解 U, S, Vt svd(W, full_matricesFalse) # 计算目标秩 rank max(1, int(min(original_shape) * rank_ratio)) # 截断 U U[:, :rank] S S[:rank] Vt Vt[:rank, :] # 创建新层 new_in_features original_shape[1] new_out_features original_shape[0] # 替换原始层 layer.weight.data torch.from_numpy(U * S).float() layer.out_features rank return U.shape[1] staticmethod def apply_to_model(model: nn.Module, rank_ratio: float 0.5): 应用到整个模型 total_reduction 0 for name, module in model.named_modules(): if isinstance(module, nn.Linear): original_size module.weight.numel() new_rank LowRankDecomposition.decompose_linear(module, rank_ratio) new_size module.in_features * new_rank new_rank * module.out_features total_reduction original_size - new_size print(fLayer {name}: {original_size} - {new_size}, 减少 {original_size - new_size}) return total_reduction # 使用 decomposer LowRankDecomposition() reduction decomposer.apply_to_model(model, rank_ratio0.5) print(f总共减少参数: {reduction})3. 性能对比3.1 压缩效果对比方法压缩比ImageNet Top-1加速比原始ResNet-501.0x76.1%1.0x剪枝(50%)2.0x75.3%1.5x量化(INT8)4.0x75.8%2.0x蒸馏可变75.5%可变剪枝量化8.0x74.6%3.0x4. 最佳实践4.1 压缩策略选择场景推荐策略理由移动端部署量化显著减小模型大小边缘设备剪枝量化综合优化实时推理知识蒸馏保证精度云端部署剪枝减少计算量4.2 注意事项# ✅ 推荐渐进式压缩 # 先剪枝再量化避免一次性压缩导致精度大幅下降 # ✅ 推荐保留关键层 # 不要对embedding层、最后一层分类器进行大幅剪枝 # ❌ 避免过度压缩 # 压缩比过高会导致精度大幅下降5. 总结模型压缩要点剪枝减少参数量和计算量量化减小模型存储和加速推理知识蒸馏用大模型指导小模型学习