蒸馏+剪枝协同压缩:面向Apple Silicon的人脸识别模型实战

发布时间:2026/9/17 4:44:53
蒸馏+剪枝协同压缩:面向Apple Silicon的人脸识别模型实战 简介本资源是一份面向深度学习初学者与毕业设计学生的模型压缩实践代码包聚焦知识蒸馏与网络剪枝两大轻量化技术解决识别模型在边缘设备或资源受限场景下的部署难题。压缩包共185个文件含79个Python源码涵盖训练、蒸馏、剪枝、模型转换等核心逻辑、60个编译缓存文件pyc、11个Git配置文件及若干日志、配置与结果记录文件如train_log、json、sh脚本等整体体积仅4.03MB结构紧凑便于快速复现。已有113人下载学习适合希望掌握PyTorch框架下模型压缩全流程的开发者——不仅提供完整可运行的蒸馏与剪枝实现还包含多数据集对比实验、Apple Silicon架构适配方案以及ArcFace等典型识别模型的轻量化改造记录目录中大量record_*和log文件直观呈现不同压缩策略的性能变化轨迹。1. 这不是“瘦身”而是“重写大脑”一个毕业设计级的模型压缩实战包为什么蒸馏剪枝必须一起用很多刚接触模型压缩的同学会误以为“剪枝就是删掉不重要的权重蒸馏就是让小模型学大模型”结果一跑就崩——精度掉5个点、推理慢了反而更卡、Apple Silicon上直接报Illegal instruction。这个名为“基于模型压缩的识别算法python源码蒸馏和剪枝.zip”的仓库本质是一个闭环验证型工程它不只提供单点技术实现而是把知识蒸馏Knowledge Distillation作为剪枝前的“认知预训练”再用结构化剪枝Structured Pruning完成通道级裁剪最后通过TorchScript ML Compute统一部署到Apple Silicon。它面向的是真实落地场景——比如边缘摄像头需在M2芯片上以80ms延迟运行人脸识别同时保持ArcFace特征余弦相似度0.72。代码里没有玩具数据集而是实测过LFW、CFP-FP、IJB-C三类benchmark没有抽象接口每个.py文件都带--teacher-path、--prune-ratio、--target-arch mps等可调参数。适合两类人一是毕设/课程设计需要交完整pipeline的学生二是想快速验证蒸馏-剪枝协同增益的算法工程师。2. 蒸馏不是“抄答案”而是构建教师-学生联合优化目标2.1 为什么必须先蒸馏再剪枝从损失函数设计讲起单纯剪枝容易破坏模型对细粒度特征的敏感性尤其在人脸识别这种依赖局部纹理如眼角褶皱、鼻翼阴影的任务中。该仓库采用两阶段联合损失第一阶段用教师模型预训练ArcFace ResNet50指导学生模型轻量ResNet18学习logits分布与中间层attention map第二阶段在此基础上进行通道剪枝。核心在distiller.py中的KDWithAttentionLoss类class KDWithAttentionLoss(nn.Module): def __init__(self, alpha0.7, beta0.3, temperature4.0): super().__init__() self.alpha alpha # logits蒸馏权重 self.beta beta # attention map蒸馏权重 self.T temperature def forward(self, student_logits, teacher_logits, student_attn, teacher_attn): # 温度缩放后的KL散度logits蒸馏 kd_loss F.kl_div( F.log_softmax(student_logits / self.T, dim1), F.softmax(teacher_logits / self.T, dim1), reductionbatchmean ) * (self.T ** 2) # 注意力图L2距离attention蒸馏 attn_loss F.mse_loss(student_attn, teacher_attn) return self.alpha * kd_loss self.beta * attn_loss提示alpha0.7并非固定值实际在train_distill.py中通过--kd-alpha 0.65动态传入。当teacher在LFW上Acc99.2%而student初始仅96.1%时调高alpha至0.75能加速logits收敛若student在CFP-FP上FARFRR1e-3表现差则需降低alpha、提升beta如0.4/0.6强制关注教师网络的注意力聚焦区域。2.1.1 教师模型如何注入学生网络仓库未使用传统torchvision.models加载而是通过arcfaceTrainLog/teacher_checkpoint.pth读取完整state_dict关键在于特征对齐层插入。在models/student_resnet.py中class StudentResNet18(nn.Module): def __init__(self, num_classes10575, embedding_size512): super().__init__() self.backbone resnet18(pretrainedFalse) # 不加载ImageNet权重 # 在layer2和layer3输出后插入attention hook self.attn_hooks [] for name, module in self.backbone.named_modules(): if layer2 in name and isinstance(module, nn.Conv2d): self.attn_hooks.append(module.register_forward_hook(self._save_attn)) if layer3 in name and isinstance(module, nn.Conv2d): self.attn_hooks.append(module.register_forward_hook(self._save_attn)) def _save_attn(self, module, input, output): # 对输出特征图做全局平均池化sigmoid生成空间注意力图 attn_map torch.sigmoid(output.mean(dim1, keepdimTrue)) # [B,1,H,W] self.current_attn attn_map此设计使学生网络在训练时能同步计算自身attention map并与teacher的对应层map比对——这比单纯logits蒸馏提升CFP-FP验证集AUC 1.8个百分点见results/distill_ablation.csv。2.2 剪枝不是“随机砍”而是基于Hessian敏感度的结构化通道裁剪2.2.1 为什么不用Magnitude-based剪枝Magnitude剪枝按权重绝对值排序在人脸识别任务中易误删判别性通道。例如某通道权重均值虽小但其梯度方差极大对应对光照变化敏感的纹理响应。该仓库采用二阶泰勒展开近似敏感度Taylor Expansion Sensitivity在pruner/taylor_pruner.py中实现def compute_taylor_sensitivity(model, dataloader, device, num_batches5): # 初始化敏感度字典 {layer_name: tensor of shape [C]} sensitivity {} model.eval() for i, (x, _) in enumerate(dataloader): if i num_batches: break x x.to(device) x.requires_grad_(True) # 前向传播获取loss此处用ArcFace loss features model(x) loss arcface_loss(features, torch.randint(0, 1000, (x.size(0),)).to(device)) # 反向传播计算梯度 loss.backward() # 遍历卷积层计算每个通道的Taylor敏感度 for name, module in model.named_modules(): if isinstance(module, nn.Conv2d) and layer in name: if name not in sensitivity: sensitivity[name] torch.zeros(module.out_channels) # 计算 ∂L/∂w * w 的绝对值和对每个通道求和 grad module.weight.grad.data # [C_out, C_in, k, k] weight module.weight.data channel_sens torch.sum(torch.abs(grad * weight), dim[1,2,3]) # [C_out] sensitivity[name] channel_sens.cpu() return {k: v / num_batches for k, v in sensitivity.items()}注意num_batches5是经验阈值。若在IJB-C上剪枝后FAR升高需增至8若剪枝后推理速度未达预期说明敏感度计算批次不足导致误判应检查dataloader是否启用pin_memoryTrue。2.2.2 结构化剪枝的具体执行流程剪枝非一次性操作而是迭代式微调Iterative Pruning Fine-tuning。prune_main.py控制全流程步骤命令示例关键参数说明1. 敏感度分析python prune_main.py --mode analyze --model-path student_kd.pth --data-path ./data/lfw--prune-ratio 0.3指定首轮剪30%通道--sensitivity-method taylor固定用泰勒法2. 生成掩码python prune_main.py --mode mask --mask-path mask_0.3.pth输出mask_0.3.pth含各层通道保留索引如layer2.0.conv1.mask: [1,0,1,1,...]3. 应用剪枝python prune_main.py --mode apply --mask-path mask_0.3.pth --model-path student_kd.pth调用torch.nn.utils.prune.custom_from_mask仅剪枝conv层bn层同步缩放4. 微调恢复python train_pruned.py --model-path pruned_student.pth --lr 1e-4 --epochs 20学习率降为蒸馏阶段的1/10避免破坏已学知识表格中第3步的“bn层同步缩放”是关键细节剪枝后BN层的running_var需按通道mask重新归一化否则torch.nn.BatchNorm2d会因方差为0触发NaN。代码在pruner/apply_pruning.py中通过_sync_bn_stats函数实现。3. Apple Silicon部署不是“加个mps”而是ML Compute与TorchScript的深度协同3.1 为什么不能直接model.to(mps)PyTorch的mps后端对算子支持有限torch.nn.functional.interpolate的modebilinear在MPS上会fallback到CPU导致GPU利用率20%torch.einsum在某些shape下报NotImplementedError。该仓库绕过这些限制采用TorchScript编译 ML Compute显式调度双路径3.1.1 TorchScript编译的三个强制约束在export_to_mps.py中模型导出前必须满足禁用动态控制流所有if/else分支需转为torch.where如原逻辑if x.size(0) 16: ...改为torch.where(x.size(0) 16, branch_a(), branch_b())固定输入shapetorch.jit.trace要求example_input torch.randn(1,3,112,112)ArcFace标准输入不可用torch.jit.script替换不兼容算子将F.interpolate替换为nn.Upsamplemodenearest并将einsum重写为torch.bmm# 替换前不兼容MPS def forward(self, x): return torch.einsum(bik,bjk-bij, x, x.transpose(1,2)) # 替换后兼容MPS def forward(self, x): x_t x.transpose(1, 2) # [B,K,D] - [B,D,K] return torch.bmm(x, x_t) # [B,K,D] [B,D,K] - [B,K,K]3.1.2 ML Compute调度器的底层绑定mps_executor.py通过ctypes直接调用Apple的MLCompute框架规避PyTorch MPS的算子限制import ctypes from pathlib import Path # 加载MLCompute动态库 mlc_lib ctypes.CDLL(/System/Library/Frameworks/MLCompute.framework/MLCompute) # 定义C函数签名 mlc_lib.MLCComputeCreateGraph.argtypes [ctypes.c_void_p] mlc_lib.MLCComputeCreateGraph.restype ctypes.c_void_p def create_mps_graph(model_path: str): # 将TorchScript模型序列化为MLC Graph graph_ptr mlc_lib.MLCComputeCreateGraph( ctypes.c_char_p(model_path.encode()) ) # 绑定MPS设备 mlc_lib.MLCComputeSetDevice(graph_ptr, mps) # 编译优化 mlc_lib.MLCComputeCompile(graph_ptr) return graph_ptr提示此部分需Xcode 14.3及macOS 13.3若ctypes.CDLL报错Library not loaded请确认/System/Library/Frameworks/MLCompute.framework存在且权限为r-xr-xr-x。3.2 性能对比剪枝蒸馏MPS部署的真实收益在M2 Ultra24核GPU上实测输入112×112 RGB图像模型配置推理延迟ms内存占用MBLFW Acc%CFP-FP AUC%ResNet50原始128.4112099.298.7ResNet18蒸馏42.148097.897.3ResNet18蒸馏剪枝0.428.629597.196.9上述模型MPS部署19.318297.096.8注意延迟数据来自time.perf_counter()在mps_executor.py的run_inference函数中测量排除数据加载时间。内存占用为psutil.Process().memory_info().rss / 1024 / 1024峰值。可见剪枝降低计算量MPS降低调度开销二者叠加产生非线性加速——延迟从42.1ms降至19.3ms提速2.2倍而非简单相加。4. 验证你的剪枝是否“伤筋动骨”三类必做的诊断性测试4.1 特征空间坍缩检测用t-SNE可视化剪枝前后embedding分布剪枝过度会导致不同ID的特征向量在嵌入空间中聚类失效。仓库提供utils/visualize_tsne.py关键步骤def plot_tsne_embeddings(model, dataloader, device, n_samples1000): model.eval() embeddings, labels [], [] with torch.no_grad(): for x, y in dataloader: if len(embeddings) n_samples: break x, y x.to(device), y.to(device) feat model(x) # ArcFace输出512维特征 embeddings.append(feat.cpu()) labels.append(y.cpu()) # 合并并降维 X torch.cat(embeddings)[:n_samples].numpy() y torch.cat(labels)[:n_samples].numpy() # t-SNE降维perplexity30适配人脸识别类间距离 tsne TSNE(n_components2, perplexity30, random_state42) X_tsne tsne.fit_transform(X) # 绘图不同颜色代表不同ID plt.scatter(X_tsne[:,0], X_tsne[:,1], cy, cmaptab20, s1) plt.colorbar() plt.title(t-SNE of Face Embeddings (Pruned vs Original)) plt.savefig(tsne_pruned.png, dpi300, bbox_inchestight)提示若剪枝后t-SNE图中出现大面积同色块某ID占据整个象限说明该ID特征被过度压缩若所有点呈均匀雾状无聚类说明判别性信息丢失。此时应回退到prune-ratio0.25并增加微调epoch。4.2 梯度流完整性检查用torch.autograd.gradcheck验证剪枝层反向传播结构化剪枝可能破坏梯度计算图。在test_gradient_flow.py中运行def test_pruned_layer_gradient(model, layer_namelayer2.0.conv1): # 获取剪枝后层 layer dict(model.named_modules())[layer_name] # 构造测试输入确保shape匹配剪枝后通道数 x torch.randn(2, layer.in_channels, 56, 56, requires_gradTrue) # 前向 out layer(x) # 构造伪标签计算loss target torch.randn_like(out) loss F.mse_loss(out, target) # 检查梯度 gradcheck_pass torch.autograd.gradcheck( lambda x: layer(x), x, eps1e-3, atol1e-3, rtol1e-2 ) print(fGradient check for {layer_name}: {PASS if gradcheck_pass else FAIL})若返回FAIL常见原因1BN层weight被错误剪枝应只剪bias2prune.custom_from_mask未设置prune.PruningMethod.PRE_TRAINING3forward中使用了torch.no_grad()包裹。4.3 Apple Silicon专属陷阱MPS张量生命周期管理MPS设备有独立内存池torch.tensor(..., devicemps)创建的张量不会自动释放。仓库在mps_executor.py中强制管理class MPSTensorManager: def __init__(self): self.cache {} def get_tensor(self, key: str, shape: tuple, dtypetorch.float32): if key not in self.cache: # 分配MPS张量并缓存 self.cache[key] torch.empty(shape, dtypedtype, devicemps) return self.cache[key] def clear_cache(self): # 显式清空缓存避免MPS内存泄漏 for tensor in self.cache.values(): tensor.data torch.empty(0, devicemps) # 触发内存回收 self.cache.clear() # 使用示例 manager MPSTensorManager() input_mps manager.get_tensor(input, (1,3,112,112)) # ... 推理 ... manager.clear_cache() # 必须在每次推理后调用提示若MPS推理多次后报RuntimeError: MPS backend out of memory90%概率是未调用clear_cache()。该方法比torch.mps.empty_cache()更可靠因后者不保证立即释放。本文还有配套的精品资源点击获取