PyTorch实现TopK剪枝：原理、代码与实战踩坑记录

最近在调一个部署项目，原来跑得好好的模型，到了目标设备上一看，体量超了、延迟也压不下去。整个压缩方案里，最先落地的倒不是那些花哨的蒸馏、量化，而是最直观的一招——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, padding=1) self.bn1 = nn.BatchNorm2d(16) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) 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_acc=98.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生成mask

PyTorch里的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()之前，是因为你要在优化器拿到梯度之前就把不该更新的位置清零。如果放后面，动量等优化器内部状态已经根据原始梯度改变了，等于没挡住。

验证时用普通的前向推理就行，不涉及训练循环。但要注意，推理时模型里如果有BatchNorm，model.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_masks=current_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(f"overall 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 问题速查表

最后整理一个速查表，方便你遇到问题时直接对照。

最后多说一句

TopK剪枝是我做模型压缩实验时第一个跑通的方案，它简单、直接、代码量少，但正因为简单，反而让你能快速建立起对模型稀疏行为的直觉。我的习惯是先从10%剪起，看一眼分层稀疏度和精度变化，再决定要不要往上加码。如果你的后续目标是真正的端侧部署加速，那我强烈建议你在TopK剪枝跑通之后，继续深入研究结构化剪枝和量化，它们才是真正让模型在硬件上跑得快的手段。而TopK剪枝作为第一步，永远不会白做——它告诉了你模型到底藏了多少可以去掉的冗余。