最近在调一个部署项目原来跑得好好的模型到了目标设备上一看体量超了、延迟也压不下去。整个压缩方案里最先落地的倒不是那些花哨的蒸馏、量化而是最直观的一招——TopK剪枝。PyTorch里实现TopK剪枝并不复杂核心思路就一句话把权重里绝对值最大的那一部分留下来其余全部置零。听起来简单但真正动手做的时候涉及到怎么选阈值、怎么不改坏训练循环、剪完之后要不要重训练这些细节如果不弄清楚很容易剪出一个“废了”的模型。这篇文章不绕弯子直接从原理讲到能跑通的完整代码再把我实际踩过的坑一并列出来。适合那些想快速给自己的模型减负、手上已有PyTorch基础、之前没系统接触过剪枝的读者。看完你会知道TopK剪枝跟阈值剪枝、随机剪枝有什么区别逐层剪和全局剪该怎么选以及最关键的——剪完之后模型精度到底靠什么拉回来。1. 剪枝到底在干什么TopK剪枝的来龙去脉很多做模型压缩的初学者上手剪枝心态都是“把没用的参数删掉”。这个想法没错但不够精确。剪枝真正利用的是参数分布里的长尾效应。一个训练充分的网络它的权重分布往往两边极不对称一小部分权重绝对值很大它们承载着网络的主要表达能力而大量权重绝对值极小接近零对最终输出的贡献微乎其微。TopK剪枝就是在这种分布上做文章——按绝对值排序保留最大的K个剩下的强制归零。你可能会问为什么不直接选一个绝对值阈值小于阈值的归零这不更简单吗确实有这种“阈值剪枝”实现起来比TopK还省事。但阈值剪枝有个问题不同层的参数尺度差异可能很大。某些层的权重普遍很小按全局阈值剪这一层会被整个端掉某些层权重普遍偏大按全局阈值剪这一层又基本没剪到。TopK的好处在于它按“数量”来控制每层保留多少个权重是显式的行为可预期。虽然PyTorch里有现成的剪枝工具但TopK这种方式让你对模型里每个位置的去留一目了然。TopK剪枝还有个别名叫“非结构化剪枝”因为它是在单个权重级别上做取舍剪完的权重矩阵里稀疏零不是按行或按列排列的而是星星点点散在各处。与之相对的是结构化剪枝比如整行、整列、整个卷积通道一起剪掉。结构化剪枝对硬件友好能直接换成更小的矩阵非结构化剪枝则更像一个“纯度”很高的实验——它用置零的方式验证模型里到底有多少参数是冗余的。实际上很多后续的稀疏化部署方案第一步就是用TopK剪枝先摸一遍模型的水深。1.1 稀疏化的本质为什么剪掉权重反而能用要理解为什么剪掉一部分权重模型还能工作得从参数冗余说起。神经网络在训练过程中会不断调整权重来拟合数据但最终学到的内部表示往往存在大量重叠和冗余——同一个功能可能被好几个权重组合分别覆盖。这跟拍照时同一场景拍一百张照片有些浪费的道理差不多。只要去掉一部分重复度高的“照片”剩下的仍然能还原出主体内容。当然不是任何一层都能随便剪。通常卷积层和全连接层的权重是剪枝的主要对象因为它们的参数量占了大头。而BatchNorm层的缩放参数weight和bias通常不参与这种逐权重的TopK剪枝因为它们是一维向量每个值直接对应一个输出通道的缩放系数。如果你把某个BatchNorm的缩放系数置零等于直接抹掉了这个通道的整个输出信号性质上已经接近结构化剪枝了而且很容易让后面的层算出的方差出问题。我在代码里会专门跳过这类参数。剪完之后模型里会多出一批零权重。这些零权重在前向计算里不会对输出产生任何贡献。但要注意它们占的内存还是那么多浮点运算量也一分不少。也就是说TopK剪枝在最原始状态下并不会直接让模型跑得更快真正的收益在后续——要么用稀疏张量格式存储要么配合专用的稀疏算子来跳过这些零。先把这一步做对再去想加速顺序不能反。1.2 逐层TopK和全局TopK怎么选TopK剪枝按照作用范围又分成“逐层TopK”和“全局TopK”。逐层TopK的做法很直白每一层独立跑一次排序保留该层权重的TopK。比如你设置剪枝比例50%那么每一层都保留自己一半的权重。好处是公平每一层都有存活名额网络结构不容易因为某一层被剪秃而崩塌。全局TopK则是把所有层的权重全部倒进同一个池子里排序按全局TopK来确定一个统一的保留阈值。这样做的好处是“按需分配”——那些本身就有大量小权重的层会被多剪一些而权重普遍较大的重要层只被轻轻动一下。听起来更灵活但风险也很明显遇到某个层参数的绝对值整体偏低时它可能被剪掉95%直接把这一层的表达能力剪没了后续层拿到的东西也会变得很稀薄。我的建议是第一次做剪枝实验时从逐层TopK开始。因为它最容易定位问题——精度掉得厉害你能迅速知道是哪一层的比例设得太狠全局TopK则适合已经有一定剪枝经验、希望压榨更高稀疏度的人。两者代码差异不大后面实现部分我都会给出。2. 动手之前环境准备和基线记录技术在跑起来之前先把地基打好。我做剪枝实验通常固定一套环境不太喜欢什么新潮用什么实验可复现比什么都重要。2.1 最小可复现的模型与数据这篇文章里的示例我建议直接用一个两层的小卷积网络跑在公开基准数据集上。网络结构越简单剪枝带来的变化越容易观察。代码如下import torch import torch.nn as nn import torch.nn.functional as F class TinyNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 16, 3, padding1) self.bn1 nn.BatchNorm2d(16) self.conv2 nn.Conv2d(16, 32, 3, padding1) self.bn2 nn.BatchNorm2d(32) self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Linear(32, 10) def forward(self, x): x F.relu(self.bn1(self.conv1(x))) x F.relu(self.bn2(self.conv2(x))) x self.pool(x).flatten(1) return self.fc(x)后面所有剪枝代码都围绕这个结构展开但它同样适用于更深的网络。数据方面用公开的手写数字基准集就行输入是1通道正好跟上面小网络对上。注意环境里PyTorch版本最好在2.x以上更低版本里的torch.topk行为也一致只是部分API细节略有差异不影响整体逻辑。2.2 基线准确率为什么必须留档剪枝实验有一个特别容易被新手忽略的动作先把未剪枝模型的准确率完整记下来。不是只记一个数而是包括训练epoch数、优化器、学习率、batch size这整套配置。因为剪枝完之后你要重训练重训练的策略必须参考原来的训练配置来做调整。我自己的做法是训练一个固定的“基准模型”脚本每次实验前自动打印一份类似这样的报告baseline_acc98.51% config{epochs:10, lr:0.01, optimizer:sgd, batch_size:64}剪枝后对比时不要只看一个两个点要看完整趋势。比如剪50%后重训练3个epoch能达到98.2%那说明这50%确实可以省如果剪50%后无论怎么重训练都只有96%那就要考虑是不是剪到了“命脉”上。没有基线所有判断都会变成拍脑袋后面排查问题也会很被动。3. 核心实现PyTorch里TopK剪枝的两种写法接下来进入重头戏。我准备了两种TopK剪枝的实现方式它们各自适合不同场景你可以在自己的项目里按需选择。3.1 按层独立剪枝torch.topk生成maskPyTorch里的torch.topk可以直接返回张量中最大的K个值及其索引但直接用它去构造新权重会打断和原模型参数的关联。更稳的做法是先通过topk拿到临界阈值再生成一个布尔掩码mask然后用这个掩码去乘法更新参数。这样的好处是后续如果你想保留“原始权重 掩码”的存储方式也非常自然。import torch def per_layer_topk_mask(param, ratio): keep_ratio 1.0 - ratio num_total param.numel() num_keep max(int(num_total * keep_ratio), 1) flat_abs param.abs().flatten() values, _ torch.topk(flat_abs, num_keep) threshold values[-1] mask param.abs() threshold return mask这里有个细节值得说torch.topk返回的values已经是按绝对值从大到小排好的取最后一个元素就是“保留下来的第K大绝对值”。用这个阈值去和每个权重比较凡是绝对值大于等于这个阈值的都放进掩码里。这样等于一口气把权重做了一遍二分效率比逐个循环高得多。但要注意一种边界情况如果参数里存在大量绝对值相等的值比如权重是0.0的很多而阈值恰好也是0.0那么生成的mask里为True的数量会略多于num_keep。这个偏差在稀疏度统计时可能带来1%以内的误差一般可以接受。如果对精确的稀疏度有强迫症可以在mask生成之后再随机挑一些True位置改成False凑到精确数量。但说实话模型训练中这种误差不会造成可感知的影响我做实验时通常直接忽略。拿到掩码之后把掩码存到一个字典里等会儿用来对模型做统一处理def build_layerwise_masks(model, ratio): masks {} for name, param in model.named_parameters(): if param.dim() 2: continue if weight not in name: continue masks[name] per_layer_topk_mask(param, ratio) # 每一层保留同样比例 return masks跳过一维参数的原因前面说过——它们往往是BatchNorm的缩放或偏置剪了会直接破坏通道结构。另外weight not in name这个判断看起来有点重复但很实用它能防止误伤其他名字里带weight但语义不同的参数。3.2 全局统一剪枝二次遍历生成mask如果你决定走全局TopK路线代码思路稍作调整先把所有符合条件的权重绝对值拼起来用torch.cat合成一个大一维张量对这个大张量做一次torch.topk拿到全局阈值再回头给每一层生成掩码。def build_global_masks(model, ratio): keep_ratio 1.0 - ratio candidates [ param.abs().flatten().detach() for name, param in model.named_parameters() if weight in name and param.dim() 2 ] all_abs torch.cat(candidates) num_keep max(int(all_abs.numel() * keep_ratio), 1) values, _ torch.topk(all_abs, num_keep) threshold values[-1] masks {} for name, param in model.named_parameters(): if weight in name and param.dim() 2: masks[name] param.abs() threshold return masks注意一个细节all_abs是从不同层拼接来的每层参数取值范围差异很大拼接后排序没问题。但这里千万别用.detach()省略否则你等于构建了一个大计算图去跟踪所有参数的梯度既占内存又可能在后续反向传播时埋下意料之外的坑。我们只是要一个阈值不需要梯度流经这里。全局TopK剪枝下来有些层的mask里可能几乎全是False这层等于被移除得所剩无几。如果发现这种情况说明全局剪枝比例设得超出了模型能承受的范围。一旦某个层存活权重低于原来10%精度很可能会迅速崩掉。这时候就得回到逐层方案或者对该层单独设一个保底比例。3.3 训练循环的正确姿势梯度mask和权重mask有了mask怎么把它用进训练循环是整个剪枝过程里最关键的一步。很多人以为只需要把权重乘上mask就行结果训练几轮之后发现精度回升了一些但稀疏度悄悄变少了——那些原本被置零的权重又活了。原因很简单optimizer.step()更新参数时每轮都会执行类似param - lr * grad的操作即便grad在置零位置上是0如果优化器带动量或者weight decay它仍可能在这些位置上产生非零更新。所以正确做法是双管齐下。第一割断梯度。反向传播之后、调用optimizer之前把梯度也乘上mask这样任何优化器收到的梯度都是“只针对存活权重”的。第二锁定权重。optimizer.step()之后强制把权重里mask为False的位置再次归零。操作顺序是先反向梯度、再mask梯度、再step、最后mask权重。配合代码看更清楚def train_one_epoch(model, masks, loader, optimizer, criterion): model.train() total_loss 0.0 for images, targets in loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, targets) loss.backward() with torch.no_grad(): for name, param in model.named_parameters(): if name in masks: param.grad.mul_(masks[name]) optimizer.step() with torch.no_grad(): for name, param in model.named_parameters(): if name in masks: param.data.mul_(masks[name]) total_loss loss.item() return total_loss / len(loader)这里把param.grad.mul_(masks[name])放在optimizer.step()之前是因为你要在优化器拿到梯度之前就把不该更新的位置清零。如果放后面动量等优化器内部状态已经根据原始梯度改变了等于没挡住。验证时用普通的前向推理就行不涉及训练循环。但要注意推理时模型里如果有BatchNormmodel.eval()必须调用否则BatchNorm会继续用batch统计量算均值和方差在稀疏权重上产生不一致的结果。我见过好几个人因为忘了调eval()剪枝后验证集上指标忽高忽低以为剪枝出了问题其实是BatchNorm状态没切过来。4. 剪完不是结束重训练和稀疏统计代码能跑通只是拿到了一个“看起来很稀疏”的模型。真正决定剪枝方案能不能用的是后续重训练和统计环节。这里面的门道一点不比剪枝本身少。4.1 剪枝比例怎么定从10%到90%剪枝比例并没有万能值但有一个通用的阶梯经验。10%~30%的剪枝属于“几乎无损”区间通常不需要特别的重训练计划稍微微调几个epoch就能恢复到接近基线水平。50%是个分水岭剪完精度会有一个可感知的下降必须重训练而且重训练的学习率不能按原来那么大。70%以上属于高压缩区间一次性剪到这个比例模型大概率直接崩掉这时候要改用迭代式剪枝。学习率的选择上我通常把初始学习率降为原来训练的1/5到1/10。原因在于剪枝后的模型相当于在一个新的损失地形上原有的大学习率很容易让剩下的少量权重剧烈震荡反而找不回原来的低损失区域。用较小学习率慢慢微调是稀疏模型重训练里最稳的策略。比如原训练用0.01剪枝后重训练推荐从0.002开始跑5个epoch再降低一半。4.2 一次性剪枝 vs 迭代式剪枝你可能以为剪枝是“一步到位”的其实实践中用得更多的反而是迭代式剪枝。所谓迭代式不是一口气剪到90%而是先把模型剪掉10%重训练让它恢复精度然后在这次剪完的基础上再剪掉10%重训练再剪……直到目标稀疏度。整个过程画出来像走台阶每一步的精度损失都很小最终累积出很高的剪枝比例。我对比过同一模型在同一种子下一次性剪80%和迭代8次剪到80%迭代式最终精度能高出2到3个百分点甚至更多。原因也好理解每剪一步后剩余权重都有机会重新适应被破坏的分布而一次剪到80%剩下那20%权重面对的是一个剧烈变形的数据分布很难独自撑起原来的表达。实现迭代式剪枝不需要改太多代码。只需要在每次剪枝时调用一次mask生成函数同时参考上一次已经置零的区域只在非零权重范围里重新做TopK。用mask乘上参数后再做排序就行。伪代码大致是current_masks None for step in range(num_iters): # 在非零权重上继续TopK逐步增加剪枝比例 prune_ratio_step 1 - (1 - final_ratio) ** (1 / num_iters) current_masks build_layerwise_masks(model, prune_ratio_step, previous_maskscurrent_masks) apply_masks(model, current_masks) for _ in range(epochs_per_step): train_one_epoch(...)每次剪枝后的重训练epoch数不需要太多3到5个epoch就好。如果每步都从头训一遍几十个epoch计算成本太高效果也不会因此更好。4.3 稀疏度统计怎么看剪完之后总得给自己一个交代模型到底稀疏到什么程度常规做法是统计全部权重中零值所占的比例。但这里容易犯一个错误直接统计参数的零值个数和按照mask统计结果很可能不一致。因为训练过程中即便你加了mask某些参数可能是极其微小的非零值比如1e-30虽然它实际起不了作用但统计时它仍然算“非零”。我习惯把mask单独拿出来统计def report_sparsity(masks): total 0 kept 0 for name, mask in masks.items(): total mask.numel() kept mask.sum().item() sparse_ratio 1.0 - kept / total print(foverall sparsity: {sparse_ratio * 100:.2f}%) return sparse_ratio这里统计的是“mask里有多少个False”也就是被剪掉的比例。同时我还会按层打印一遍稀疏度分布重点关注那些被剪得特别狠的层。如果某层稀疏度超过95%而其他层只有50%那这一层的存活权重过少后续很容易变成瓶颈。这种分层报告在排查问题时非常有用能快速定位是哪一层拖累了整体精度。5. 常见问题与排查技巧实录这一节全是我实战中踩过的坑一条一条说清楚省得你重新趟一遍。5.1 剪完精度崩了怎么办剪完精度掉得特别厉害最先检查的永远是“剪枝比例是不是太大”和“重训练是不是没跟上”。如果比例已经设在30%以内但还是崩那就要看是不是剪到了不该剪的参数。检查一下你的mask生成函数是不是把BatchNorm的weight或者bias也纳入剪枝范围了。我见过最隐蔽的错误是使用named_parameters()遍历时没有对参数名做判断结果把bn1.weight也当成普通权重给剪了精度一下子掉了十几个点。另一种情况是重训练时学习率太大模型还没恢复就已经震荡到别处去了。可以试试把学习率再降一个数量级从0.0005左右开始微调。不要怕慢剪枝后的模型参数量变小了每个epoch耗时本来就缩短多跑几个epoch成本可控。5.2 稀疏度统计不准的问题出现“我明明设置了剪50%统计出来只有40%”的情况大概率是前面提到的mask里True数量多于理论值。比如权重里有大量绝对值恰好等于阈值的元素条件把一批等于阈值的也都保留下来了。解决办法很简单生成mask后做一次数量修正多出来的True位置随机置为False直到数量精确等于num_keep。另一个常见问题出在全局TopK里如果不同层的参数数量差异悬殊大层会主导排序结果小层的权重几乎都低于阈值导致小层被剪得特别狠。这种情况统计出的全局稀疏度是正确的但分层稀疏度分布会非常畸形。遇到这个问题我建议回到逐层TopK方案或者对每个层设置一个“最低存活比例”比如无论全局阈值多高该层至少保留下10%的权重。5.3 mask丢了继续训练等于白剪这个问题特别坑你费劲剪出一个稀疏模型保存了state_dict但忘记保存mask。下次从checkpoint加载模型继续训练时训练代码里没有mask可用之前的稀疏权重在几轮更新后就会被慢慢填回非零值。剪枝成果付诸东流。解决办法是把mask和模型绑在一起保存。我通常用这样一个结构checkpoint { model: model.state_dict(), masks: {name: mask.cpu() for name, mask in masks.items()}, sparsity: report_sparsity(masks), } torch.save(checkpoint, pruned_tiny.pt)加载时先加载模型再加载mask然后把mask应用到模型上。后续只要你想继续训练就把mask传给训练函数保持整个流程闭环。5.4 剪完没有加速非结构化剪枝的局限这是很多人对TopK剪枝产生误解的地方。剪枝后模型参数量名义上少了50%但在普通PyTorch推理中运行时间几乎不变甚至可能因为稀疏度统计带来的额外操作而变得更慢。原因前面也提了非结构化剪枝产生的稀疏矩阵是散点状的默认的密集矩阵乘法不会自动跳过零值位置。要让模型真正加速有两条路。一条是使用支持稀疏张量计算的框架或库让卷积和全连接层在稀疏格式上运行另一条是先做结构化剪枝比如直接把某个卷积通道整体裁减掉这样网络层本身变得更窄推理算力需求实打实下降。TopK剪枝在工程实践里的定位更多是拿来评估“模型有没有可压的空间”探明潜力之后再决定下一步上什么手段。5.5 问题速查表最后整理一个速查表方便你遇到问题时直接对照。现象可能原因解决办法剪枝后精度暴跌剪枝比例过大/误剪了BN和bias降低比例、检查mask生成逻辑稀疏度低于预期阈值处存在大量相等值mask生成后做精确数量修正训练几轮后稀疏度下降没有在step后强制mask权重在optimizer.step后执行param.data.mul_(mask)全局剪枝后某些层被剪秃大层主导了全局阈值切回逐层方案或设置最低存活比例推理速度没有提升非结构化稀疏在密集运算下无效换稀疏算子或改用结构化剪枝加载后继续训练稀疏度丢失mask没有保存把mask写进checkpoint文件最后多说一句TopK剪枝是我做模型压缩实验时第一个跑通的方案它简单、直接、代码量少但正因为简单反而让你能快速建立起对模型稀疏行为的直觉。我的习惯是先从10%剪起看一眼分层稀疏度和精度变化再决定要不要往上加码。如果你的后续目标是真正的端侧部署加速那我强烈建议你在TopK剪枝跑通之后继续深入研究结构化剪枝和量化它们才是真正让模型在硬件上跑得快的手段。而TopK剪枝作为第一步永远不会白做——它告诉了你模型到底藏了多少可以去掉的冗余。