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

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

1. 技术分析

1.1 模型压缩方法对比

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, method='channel', amount=0.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(sparsity=0.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}, dtype=torch.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, inplace=True) # 校准 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, inplace=False) 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}, dtype=torch.qint8, inplace=True ) 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, dim=1) soft_student = F.log_softmax(student_logits / self.temperature, dim=1) soft_loss = F.kl_div( soft_student, soft_teacher, reduction='batchmean' ) * (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, epochs=10): """训练学生模型""" 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(f"Epoch {epoch+1}/{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, temperature=4.0, alpha=0.5) optimizer = torch.optim.Adam(student.parameters(), lr=0.001) distiller.train_student(train_loader, optimizer, epochs=20) 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_matrices=False) # 计算目标秩 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(f"Layer {name}: {original_size} -> {new_size}, 减少 {original_size - new_size}") return total_reduction # 使用 decomposer = LowRankDecomposition() reduction = decomposer.apply_to_model(model, rank_ratio=0.5) print(f"总共减少参数: {reduction}")

3. 性能对比

3.1 压缩效果对比

4. 最佳实践

4.1 压缩策略选择

4.2 注意事项

# ✅ 推荐：渐进式压缩 # 先剪枝再量化，避免一次性压缩导致精度大幅下降 # ✅ 推荐：保留关键层 # 不要对embedding层、最后一层分类器进行大幅剪枝 # ❌ 避免：过度压缩 # 压缩比过高会导致精度大幅下降

5. 总结

模型压缩要点：

1. 剪枝：减少参数量和计算量
2. 量化：减小模型存储和加速推理
3. 知识蒸馏：用大模型指导小模型学习