模型压缩实战：蒸馏与剪枝源码解析及边缘部署优化-开发者社区

简介：这份资源是面向毕业设计与模型压缩入门者的Python代码仓库，聚焦基于知识蒸馏与剪枝的识别算法实现，适合具备一定深度学习基础、需要完成相关课题或复现压缩实验的学生与开发者。压缩包共185个文件，约4.03MB，以79个py源码文件为核心，辅以60个pyc编译文件、若干txt说明、json配置、sh脚本及训练日志与记录文件，覆盖模型训练、剪枝、蒸馏与结果记录等环节。内容涉及知识蒸馏、模型剪枝、不同数据集上的模型对比，以及将模型转换为Apple Silicon架构等实践方向，可帮助读者理解压缩流程、对照实验配置并复用训练脚本。目前已有113人学习，适合作为毕设参考或模型压缩练手项目。

1. 模型压缩识别算法源码包：蒸馏加剪枝到底能压到什么程度

一个训练集上 98% 准确率的识别模型，部署到边缘设备上推理一次要 300 毫秒，内存占用 200MB 以上，这种场景做视觉识别的工程师基本都遇到过。模型压缩要解决的就是这个问题：在不显著掉点的前提下，把模型体积和推理耗时压下来。蒸馏和剪枝是两条最主流的路线，前者让小模型学大模型的输出分布，后者直接砍掉冗余的通道或权重。这个源码包把两条路线都做成了可跑的 Python 实现，适合已经有一个能用的识别模型、想把它塞进更小硬件里的从业者。读完你应该能判断自己的场景该走哪条路、参数怎么设、哪里容易翻车。

2. 蒸馏和剪枝的选型逻辑：先搞清楚你的瓶颈在哪

2.1 蒸馏适合什么场景，剪枝适合什么场景

蒸馏的本质是知识迁移。你有一个大模型（教师），它的 softmax 输出不只是「这个样本属于类别 A」，而是「属于 A 的概率 0.85，属于 B 的概率 0.12，属于 C 的概率 0.03」。后面这些信息叫暗知识（dark knowledge），它编码了类别之间的相似性关系。小模型（学生）直接学硬标签只能学到「A 是对的」，但学教师的软输出能学到「A 和 B 比较像，和 C 差很远」。这就是为什么蒸馏出来的小模型往往比直接用硬标签训练的同结构模型高 1 到 3 个点。

蒸馏适合的场景很明确：你手头已经有一个精度不错但跑不动的大模型，同时愿意接受一个结构更小的学生网络。学生网络的结构可以自己设计，也可以直接用现成的轻量骨干（比如 MobileNet 系列、ShuffleNet 系列）。蒸馏不改变学生网络的结构，它改变的是训练信号。

剪枝的逻辑完全不同。它假设训练好的网络里有大量冗余参数，把这些参数去掉，精度不会掉太多。剪枝分两类：非结构化剪枝是把单个权重置零，产生稀疏矩阵，理论上能压缩存储，但实际推理加速需要硬件和推理引擎支持稀疏计算，否则加速效果很有限；结构化剪枝是直接砍掉整个卷积核或通道，产生的是稠密的小模型，通用硬件上就能加速。源码包里两种都实现了，但如果你追求实际推理加速，优先看结构化剪枝那条路。

选型判断可以用一个简单规则：如果你的瓶颈是「没有小模型可用，但有大模型和训练数据」，走蒸馏；如果瓶颈是「已经有一个结构还行但参数冗余的模型」，走剪枝；如果两者都想要，先剪枝再蒸馏，或者交替做。

2.2 源码包的整体结构和依赖环境

拿到一个压缩源码包，第一件事不是跑训练，而是把目录结构和依赖关系摸清楚。常见的组织方式是按方法分目录，每个方法下面有独立的模型定义、训练脚本和配置文件。依赖方面，PyTorch 是主力框架，版本建议 1.10 以上，因为要用到torch.nn.utils.prune模块和一些较新的算子。其他依赖包括 numpy、tqdm、tensorboard（可选，看训练日志用）。

环境配置这一步很多人翻车。Python 版本建议 3.8 到 3.10，太新的版本某些 PyTorch 轮子还没跟上。如果你用 conda，创建一个独立环境是最稳妥的做法：

conda create -n model_compress python=3.9 conda activate model_compress pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install numpy tqdm tensorboard

如果你用 venv 而不是 conda，逻辑一样，先隔离环境再装包。CUDA 版本根据你显卡驱动选，cu118对应 CUDA 11.8，老卡可能要用cu113或cu102。装完之后用一行命令验证：

import torch print(torch.__version__) print(torch.cuda.is_available())

如果第二行输出False，说明 CUDA 没配好，后面训练会退回 CPU，速度差几十倍。这是第一个要排查的点。

2.3 蒸馏的核心实现：损失函数和温度参数

蒸馏的损失函数是两部分加权：一部分是学生输出和教师软输出的 KL 散度，另一部分是学生输出和真实标签的交叉熵。温度参数 T 控制软输出的平滑程度，T 越大，概率分布越平滑，暗知识越丰富，但太大会让分布接近均匀，反而丢失信息。常见取值是 3 到 10 之间，分类任务上 T=4 或 T=5 是比较稳的起点。

下面是一个蒸馏损失的核心实现，你可以直接对照源码包里的对应文件看：

import torch import torch.nn as nn import torch.nn.functional as F class DistillationLoss(nn.Module): def __init__(self, temperature=4.0, alpha=0.7): super().__init__() self.T = temperature self.alpha = alpha # 蒸馏损失权重 self.ce = nn.CrossEntropyLoss() def forward(self, student_logits, teacher_logits, labels): # 软标签损失：KL 散度，注意要乘 T^2 保持梯度量级 soft_loss = F.kl_div( F.log_softmax(student_logits / self.T, dim=1), F.softmax(teacher_logits / self.T, dim=1), reduction='batchmean' ) * (self.T ** 2) # 硬标签损失 hard_loss = self.ce(student_logits, labels) # 加权求和 return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

这段代码里有两个容易忽略的点。第一，soft_loss乘了T^2，原因是 softmax 除以 T 之后梯度会缩小 T^2 倍，不补回来蒸馏损失的梯度会太小，训练不动。第二，alpha控制软硬损失的比例，经验值在 0.5 到 0.9 之间，学生和教师差距越大，alpha 可以适当调大，让学生更多依赖教师的软信号。

教师模型在蒸馏训练过程中要冻结参数，并且切到 eval 模式。如果忘了切 eval，BatchNorm 层会用当前 batch 的统计量，导致教师输出不稳定，学生学到的信号有噪声。这个坑很隐蔽，因为 loss 不会报错，只是最终精度差一截。

2.4 剪枝的核心实现：结构化通道剪枝的流程

结构化剪枝的流程比蒸馏多几个步骤：先训练一个基准模型，然后评估每个通道的重要性，按重要性排序去掉最不重要的通道，最后对剪枝后的模型做微调恢复精度。重要性评估的准则有很多种，L1 norm 是最简单也最常用的：一个卷积核的权重绝对值之和越小，说明这个核的输出越小，越不重要。

import torch.nn.utils.prune as prune def prune_conv_l1(model, amount=0.3): """对模型中所有 Conv2d 层做 L1 结构化剪枝""" for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): prune.ln_structured( module, name='weight', amount=amount, n=1, dim=0 ) return model

这里dim=0表示按输出通道维度剪，n=1表示用 L1 norm。amount=0.3表示剪掉 30% 的通道。注意 PyTorch 的prune模块默认是「重参数化」方式，它给权重加了一个 mask，并没有真正把通道删掉，推理时计算量没变。要真正加速，需要调用prune.remove把 mask 固化，然后手动重建一个更窄的模型，把保留的通道权重拷过去。源码包里通常会有compress_model或rebuild_model这样的函数来做这件事，你重点看那部分。

剪枝率不能一次设太高。经验做法是迭代剪枝：每次剪 10% 到 20%，微调几个 epoch，再剪下一轮。一次性剪 50% 以上，精度基本会崩，微调也救不回来。这个血泪经验在结构化剪枝上尤其明显。

3. 从零跑通蒸馏训练：数据准备、教师加载和学生配置

3.1 数据加载和教师模型加载的注意事项

数据部分用标准的torchvision.datasets.ImageFolder或自定义 Dataset 都行，关键是训练集和验证集的划分要固定随机种子，否则每次跑出来的精度波动会让你误以为是蒸馏参数的问题。教师模型的加载有两种方式：如果教师是 torchvision 自带的预训练模型，直接models.resnet50(pretrained=True)就行；如果是自己训练的 checkpoint，用torch.load加载 state_dict，注意map_location要设对，否则 GPU 上存的模型在 CPU 环境加载会报错。

import torch import torchvision.models as models # 加载教师模型 teacher = models.resnet50(pretrained=False) teacher.load_state_dict(torch.load('teacher_best.pth', map_location='cpu')) teacher.eval() for param in teacher.parameters(): param.requires_grad = False # 学生模型：用更轻的骨干 student = models.resnet18(pretrained=False)

教师模型一定要先eval()再冻结梯度。顺序反了不影响功能，但养成先 eval 的习惯能避免很多玄学问题。学生模型的结构选择上，分类任务用 ResNet18 或 MobileNetV2 都行，检测任务学生骨干要和检测头匹配，不能随便换。

3.2 训练循环里蒸馏 loss 的接入方式

训练循环和普通训练几乎一样，唯一区别是每个 batch 要同时跑教师和学生，然后把两个 logits 一起送进蒸馏损失函数。下面是一个最小训练循环的骨架：

optimizer = torch.optim.SGD(student.parameters(), lr=0.01, momentum=0.9) criterion = DistillationLoss(temperature=4.0, alpha=0.7) for epoch in range(num_epochs): student.train() for images, labels in train_loader: images, labels = images.cuda(), labels.cuda() with torch.no_grad(): teacher_logits = teacher(images) student_logits = student(images) loss = criterion(student_logits, teacher_logits, labels) optimizer.zero_grad() loss.backward() optimizer.step()

教师前向要包在torch.no_grad()里，否则会占额外显存，batch size 大一点就 OOM。学习率方面，蒸馏训练的学习率可以比从头训练稍小，因为学生有教师引导，不需要太大的步长去探索。常见设置是 0.01 到 0.05 之间，配合 cosine 衰减。

3.3 蒸馏效果验证：看哪些指标，怎么对比

验证不能只看学生模型的 top-1 准确率。你至少要对比三组数字：学生从头训练（无蒸馏）的精度、学生蒸馏后的精度、教师模型的精度。如果蒸馏后的学生比无蒸馏学生高不到 0.5 个点，说明蒸馏没起作用，要检查温度、alpha 和教师输出是否正常。一个快速检查方法是打印教师 softmax 输出的熵，如果熵接近 0（分布太尖锐），说明温度太低，暗知识没提取出来；如果熵接近 log(类别数)（分布太均匀），说明温度太高。

另外，蒸馏训练收敛通常比普通训练慢，因为软标签的梯度信号更平滑。不要跑几个 epoch 看 loss 没降就放弃，至少跑完完整的学习率衰减周期再判断。

4. 剪枝实操：重要性评估、迭代剪枝和微调恢复

4.1 用 L1 norm 做通道重要性排序

在动手剪之前，先把每一层卷积核的 L1 norm 算出来看看分布。如果某一层的 norm 分布很集中，说明这层冗余度高，可以多剪；如果分布很分散，说明每个通道都在起作用，剪多了会掉点。

def compute_channel_importance(model): importance = {} for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): # 每个输出通道的 L1 norm norm = module.weight.data.abs().sum(dim=(1, 2, 3)) importance[name] = norm.cpu().numpy() return importance

拿到 importance 之后，可以画个直方图或者按层打印 min/max/mean。经验上，靠近输入的层和靠近输出的层对剪枝更敏感，中间层冗余度更高。所以剪枝率可以分层设置：浅层剪 10% 到 20%，中间层剪 30% 到 40%，深层剪 20% 到 30%。

4.2 迭代剪枝的完整流程和微调策略

迭代剪枝的伪代码逻辑是这样的：训练基准模型 → 评估重要性 → 剪一层或几层 → 微调 → 重复直到达到目标压缩率。微调的学习率要比初始训练小一个量级，通常用 0.001 左右，跑 10 到 20 个 epoch。微调数据用全部训练集，不要只用子集，否则恢复不充分。

def iterative_prune(model, train_loader, target_sparsity=0.5, step=0.1): current_sparsity = 0.0 while current_sparsity < target_sparsity: # 剪一步 model = prune_conv_l1(model, amount=step) # 微调 fine_tune(model, train_loader, epochs=10, lr=0.001) current_sparsity += step print(f"Sparsity: {current_sparsity:.1f}, Acc: {evaluate(model):.4f}") return model

每次剪完必须微调，不能连续剪多次再一起微调，那样精度掉下去就回不来了。微调的时候建议用比初始训练更小的学习率加上 warmup，让模型慢慢适应结构变化。

4.3 剪枝后模型重建：从 mask 到真正的窄模型

前面提过，PyTorch 的prune模块只是加 mask，不改变实际计算图。要真正加速，必须重建模型。重建的逻辑是：对每一层，找出被保留的通道索引，创建一个输出通道数等于保留数量的新 Conv2d，把对应权重拷过去。如果下一层的输入通道依赖上一层的输出通道，下一层的输入也要同步裁剪。这就是为什么结构化剪枝比非结构化剪枝麻烦，它涉及层与层之间的依赖关系。

def rebuild_conv(conv, keep_indices): """根据保留索引重建卷积层""" out_channels = len(keep_indices) in_channels = conv.in_channels new_conv = torch.nn.Conv2d( in_channels, out_channels, kernel_size=conv.kernel_size, stride=conv.stride, padding=conv.padding, bias=conv.bias is not None ) new_conv.weight.data = conv.weight.data[keep_indices].clone() if conv.bias is not None: new_conv.bias.data = conv.bias.data[keep_indices].clone() return new_conv

重建之后要跑一遍验证集确认精度和剪枝前一致（或只差零点几个点），然后再做后续微调。如果重建后精度暴跌，大概率是通道索引对错了，或者 BatchNorm 的 running_mean/running_var 没有同步裁剪。BatchNorm 的裁剪经常被忽略，但它和卷积输出通道是一一对应的，卷积剪了哪些通道，BN 就要剪哪些。

5. 避坑与排查：蒸馏剪枝里最容易翻车的五个地方

5.1 蒸馏 loss 不下降，学生精度和从头训练一样

现象：训练日志里 loss 在降，但验证集精度和不用蒸馏的学生模型几乎一样，甚至更低。

原因：最常见的是教师模型没有切 eval 模式，BatchNorm 统计量在训练中不断变化，教师输出不稳定。其次是温度参数设得太小（比如 T=1），软标签退化成硬标签，暗知识完全丢失。

解决：确认教师模型在蒸馏前调用了eval()并且冻结了梯度。温度从 T=4 开始试，观察教师输出的熵是否在合理范围。如果还不行，检查 alpha 是不是设得太小，软损失被硬损失淹没了。

5.2 剪枝后模型推理速度没变

现象：剪了 40% 的通道，验证精度也还行，但推理耗时和剪枝前一样。

原因：用了 PyTorch 的prune模块但没有调用prune.remove和重建模型，实际计算图没变，只是权重被 mask 了。另一种可能是用了非结构化剪枝，产生稀疏矩阵但推理引擎不支持稀疏加速。

解决：结构化剪枝后必须重建模型，把 mask 变成真实的窄层。非结构化剪枝如果目标平台不支持稀疏计算，就不要指望推理加速，它只省存储。

5.3 微调后精度恢复不到剪枝前水平

现象：剪枝率 30%，微调 20 个 epoch，精度还是比基准低 3 个点以上。

原因：剪枝率一次设太高，或者微调学习率太大导致模型在窄结构上震荡。也可能是剪枝时把关键层剪太狠了，比如第一层或最后一层。

解决：降低单次剪枝率，改成每次 10% 迭代剪。微调学习率降到初始训练的十分之一，加 warmup。检查各层剪枝率是否均匀，浅层和深层少剪。

5.4 蒸馏训练显存不够，batch size 上不去

现象：教师和学生同时前向，显存占用是单模型的两倍多，batch size 只能设很小，训练不稳定。

原因：教师前向没有包在torch.no_grad()里，PyTorch 为教师也建了计算图，显存直接翻倍。

解决：教师前向必须用with torch.no_grad():包住。如果显存还是紧张，可以把教师输出提前算好存成文件，训练时直接读 logits，这样训练阶段只需要跑学生。

5.5 剪枝后 BatchNorm 统计量不匹配导致精度崩

现象：重建模型后精度从 95% 掉到 60% 多，但权重明明是对应拷贝的。

原因：卷积通道剪了，但后面的 BatchNorm 没有同步裁剪，通道数对不上，或者 running_mean/running_var 还是旧的。

解决：重建卷积的同时重建对应的 BatchNorm，把保留通道的 running_mean、running_var、weight、bias 一起拷过去。重建后在训练集上跑几百个 batch 让 BN 重新统计，再做微调。

6. 进阶技巧：蒸馏和剪枝交替做，以及一个验证压缩收益的硬指标

单独做蒸馏或剪枝，压缩率到 2 到 3 倍之后就会遇到瓶颈。继续剪，精度崩；继续蒸，学生容量不够学不动。这时候可以交替做：先剪枝得到一个中间模型，用它当教师去蒸馏一个更小的学生，学生训练完再剪一轮。这个流程能把压缩率推到 5 倍以上，但每一步都要验证精度，不能跳步。

一个具体的交替流程是这样的：基准模型剪枝 30% 得到模型 A，微调恢复精度；用模型 A 当教师，蒸馏一个通道数减半的学生模型 B；模型 B 再剪枝 20%，微调。最终模型 B 的体积大概是基准的 25% 到 30%，精度通常能保持在基准的 95% 以上。这个数字不是绝对的，取决于你的任务难度和基准模型冗余度。

验证压缩收益不能只看参数量。参数量少不代表推理快，因为推理速度还受内存访问、算子实现、硬件并行度影响。真正要看的指标是：在目标硬件上的单次推理延迟（毫秒）、峰值内存占用（MB）、以及精度。这三个数字放在一起才能判断压缩方案值不值得上。我一般会做一个表格，把基准模型、蒸馏模型、剪枝模型、交替压缩模型的这三项指标列出来，一目了然。