企业官网建设流程全解析

模型压缩这块，我上手最勤的一招就是 PyTorch 框架里的 topk 剪枝。一句话说明它在干什么：把权重张量按绝对值大小排序，选一个阈值，小于阈值的直接置成 0，再用一套 mask 机制让梯度更新时绕开这些被剪位置，之后继续训练把精度捞回来。

这技术属于细粒度非结构化剪枝，它的优势是压缩天花板高、对模型结构零破坏，缺点是推理阶段需要稀疏算子才能直接兑现加速。无论是刚接触模型压缩的新同学，还是已经被部署体积卡住的老手，都可以把本文当成一条可直接复现的完整路径：从原理选型到完整代码，再到常见坑位排查。

1. 剪枝到底在剪什么：从“去弱连接”说起

1.1 为什么绝对值小的权重可以移除

训练好的网络里，权重绝对值分布通常是长尾的，绝大多数权重都挤在 0 附近。这背后有 L2 正则化的作用：权重大幅偏大时惩罚也大，训练过程自然会把这些不重要的连接推向小值。这些“小值连接”对输出贡献很弱，当 batch 内数据被平均之后，它们几乎不承载有效信息。把它们置 0，对 loss 的影响可以看作一次微小扰动，而神经网络在收敛点附近的损失面往往相当平坦，这种扰动很快就能被后续微调拉回来。

这就是为什么剪掉 50%~80% 的权重通常不掉点的原因。但有个边界前提：模型要训练充分。一个欠拟合的模型权重分布本来就没稳定下来，你按绝对值剪，剪掉的很可能是还在“发育期”的连接，后续恢复就会很吃力。所以实操上我从来不建议拿没训好的模型直接做剪枝。

1.2 topk 剪枝的定位：非结构化细粒度剪枝

按剪枝粒度划分，权重级剪枝也叫非结构化剪枝，是把单个权重元素置 0；通道级剪枝是把整个卷积核或输出通道删掉；块级剪枝则要考虑硬件对齐，在局部连续区域内做删除。topk 天然属于权重级，因为它只按数值大小决定单个元素去留，完全不要求“被删单元结构规整”。

为什么优先做细粒度？因为结构约束越少，压缩率的天花板越高。通道剪枝要考虑下一层输入通道数必须匹配，剪一层往往牵动全局，推到 90% 以上稀疏度几乎不可能；细粒度剪枝理论上能把参数量压到 5% 甚至 1%，代价是推理阶段需要稀疏算子的支持。简单类比：通道剪枝像拆承重墙，必须重新设计户型；topk 剪枝像把家里杂物的边角去掉，户型完全不变，但仓库空间立刻省出来了。

1.3 我为什么把 topk 剪枝当成第一板斧

如果让我给“快速交付一个压缩 Demo”选题，我每次都推 topk 剪枝，理由有三个。

实现成本低，只需要 torch.topk 和原生张量操作就能写完核心逻辑。效果稳定性好，大多数图像任务上 90% 稀疏度经过微调仍然能保住 97% 以上的精度。模型兼容性强，与通道剪枝必须修改网络定义不同，topk 剪枝完全不动 forward 计算图结构，同一套剪枝代码可以打包成通用工具，接到任意 PyTorch 模型后面。

另外，topk 剪枝是后续一切剪枝研究的“底层砖”。理解它的 mask 机制之后，再去看结构化剪枝、彩票假设、剪枝与量化联动，都会轻松很多。我见过不少同学一上来就啃结构化剪枝论文，结果被各种通道重排和约束搞到劝退，回过头来才发现 topk 这套基本功没打好。

2. 动手前的关键设计：mask 才是真正的核心逻辑

2.1 两种实现：逐层 topk 与全局 topk

先看最简单的逐层实现，每一层各自保留绝对值最大的 K 个权重：

def layerwise_topk_prune(model, sparsity): masks = {} for name, p in model.named_parameters(): if p.dim() < 2: continue keep = max(int(p.numel() * (1 - sparsity)), 1)

values, indices = torch.topk(p.data.abs().reshape(-1), keep) mask = torch.zeros(p.numel(), dtype=torch.float32, device=p.device) mask[indices] = 1.0 masks[name] = mask.reshape_as(p) return masks

这里 indices 是对 p.data.abs().reshape(-1) 这个一维视图的索引，等价于原张量摊平后的位置。把对应位置 mask 置 1，再 reshape 回原始 shape，刚好保证了每层恰好保留 K 个权重。

全局实现则是把所有权重绝对值拼到一起，定一个统一阈值，骨架如下：

def global_topk_prune(model, sparsity): all_abs = [] prune_names = [] for name, p in model.named_parameters(): if p.dim() >= 2: all_abs.append(p.data.abs().reshape(-1)) prune_names.append(name)

all_abs = torch.cat(all_abs) keep_num = max(int(all_abs.numel() * (1 - sparsity)), 1) threshold = torch.topk(all_abs, keep_num).values[-1] masks = {} for name, p in model.named_parameters(): if name in prune_names: masks[name] = (p.data.abs() >= threshold).float() return masks

逐层实现保证每层稀疏度严格一致，但会让不重要层也保留大量参数。全局实现把所有参数混在一起排序，阈值统一，自动把“保留名额”倾斜给权重更大的层。实验上全局剪通常比逐层剪恢复得更好，这也是我在文中例子里会采用的方案。

2.2 全局剪枝为什么通常比逐层剪枝好

网络各层的冗余度差异很大。输入层附近的权重往往均值偏小，但又承担着最基础的特征提取；最后一层分类头通常权重大，直接影响 logits。如果强行每层剪相同比例，可能会在第一层留下垃圾参数，却又在分类头剪掉重要权重。全局阈值解决的是“让数据决定每层剪多少”的问题。

我一般用经费分配来类比：与其每个部门按统一比例砍预算，不如按投入产出比来砍。全局剪枝做的事情就是在全局排序之后，把预算优先分配给产出最高的参数。代价是等值边界时 mask 可能多保留一些权重，实际稀疏度会略低于目标 0.1~0.5 个百分点，这点误差在大型模型上通常可以忽略。

2.3 bias 和归一化参数为什么不进剪枝范围

代码里用 p.dim() < 2 跳过一维参数，这个判断不是随手写的。bias 数量少，每个元素都直接参与输出偏移，剪掉一个相当于给某个通道加常数偏移，损失往往是灾难性的。BN 层的 gamma 和 beta 更特殊，它们的“绝对值”本身没有可比性，因为这个值和输入分布直接相关，不能拿来做重要性排序。

工程上默认只剪二维以上的 weight 张量。但要注意一个特例：NLP 任务里的 embedding 层也是二维的，按这个规则会被剪。embedding 稀疏化会直接破坏词向量的语义表达，因此做语言模型剪枝时，我通常会显式加一个 exclude_names，把 embedding 排除在外。这件事文档里很少写清楚，实际项目里影响很大。

2.4 mask 的三种用法：梯度、权重、持久化

剪枝不是一锤子买卖，剪完之后通常要微调。微调阶段最大的隐患是优化器更新会让 0 权重重新变非零。所以 mask 要做三件事。

剪枝时把 weight.data 乘以 mask 直接置零。每步 backward 之后把 param.grad.data 乘以 mask，让这些位置的梯度彻底死掉。训练完成后保存 mask 或非零索引，否则下次加载模型，稀疏结构就完全丢了。

对应的梯度 mask 函数写法如下：

def apply_grad_mask(model, masks): for name, p in model.named_parameters(): if name in masks and p.grad is not None: p.grad.data.mul_(masks[name])

这里藏着一个非常关键的坑：如果优化器配置了 weight_decay，即便梯度被 mask 置零，optimizer.step() 内部仍然会把 weight decay 项加到更新里，0 的位置会重新获得一个小幅非零值。所以更稳妥的做法是在 optimizer.step() 之后再强制做一次 weight.data *= mask，或者微调阶段直接把 weight_decay 设为 0。我在下面的完整实现里用了“step 之后再回贴 mask”的方式，这样即使优化器偷改了权重，也能立刻纠正回来。

3. 完整跑通一套 PyTorch 实现

3.1 实验配置：小模型也能说明大问题

为了把代码长度压到最短，我用公开手写数字数据集，模型选 LeNet 风格的简单 CNN。选这个组合不是为了刷榜单，而是让整个流程在一台普通 CPU 机器上也能跑完，方便你逐步验证每个环节。

import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms

class SimpleCNN(nn.Module): definit(self): super().init() self.conv1 = nn.Conv2d(1, 6, 5, padding=2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10)

def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) return self.fc3(x)

这个 44.4k 参数的小模型训练速度足够快，适合反复实验。

3.2 训练、评估、剪枝三个核心函数

训练函数里嵌入梯度 mask 和权重 mask：

def train_one_epoch(model, loader, optimizer, masks=None): model.train() for x, y in loader: optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward()

if masks is not None: apply_grad_mask(model, masks) optimizer.step() if masks is not None: apply_weight_mask(model, masks)

def apply_weight_mask(model, masks): with torch.no_grad(): for name, p in model.named_parameters(): if name in masks: p.data.mul_(masks[name])

评估函数用 no_grad 包裹，计数正确预测：

def evaluate(model, loader): model.eval() n_correct = 0 n_total = 0 with torch.no_grad(): for x, y in loader: pred = model(x).argmax(dim=1) n_correct += (pred == y).sum().item() n_total += y.size(0) return n_correct / n_total

稀疏度统计函数，用来确认剪枝结果是否符合预期：

def calc_sparsity(model): total = 0 zeros = 0 for p in model.parameters(): if p.dim() < 2: continue total += p.numel() zeros += (p == 0).sum().item() return zeros / total

3.3 完整训练主流程

先是基线训练，训练 5 个 epoch：

def main(): torch.manual_seed(3407) device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_loader = DataLoader( datasets.MNIST("./data", train=True, download=True, transform=transform), batch_size=64, shuffle=True ) test_loader = DataLoader( datasets.MNIST("./data", train=False, transform=transform), batch_size=256 ) model = SimpleCNN().to(device) optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) for epoch in range(5): train_one_epoch(model, train_loader, optimizer) acc = evaluate(model, test_loader) print(f"epoch {epoch + 1}: {acc:.4f}")

然后对几档稀疏度分别做剪枝和微调：

sparsities = [0.5, 0.8, 0.9, 0.95] for sp in sparsities: model = SimpleCNN().to(device) optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) for epoch in range(5): train_one_epoch(model, train_loader, optimizer) baseline = evaluate(model, test_loader) masks = global_topk_prune(model, sp) apply_weight_mask(model, masks) acc_after_prune = evaluate(model, test_loader) optimizer = torch.optim.SGD(model.parameters(), lr=0.002, momentum=0.9) for epoch in range(5): train_one_epoch(model, train_loader, optimizer, masks=masks) acc_finetune = evaluate(model, test_loader) print( f"sparsity={sp:.2f} " f"baseline={baseline:.4f} " f"after_prune={acc_after_prune:.4f} " f"finetune={acc_finetune:.4f}" )

注意微调阶段学习率是基线的五分之一，这个缩放不是拍脑袋定的。剪枝在权重上引入了一次硬扰动，此时再用大步长更新容易把稀疏结构冲散，小步快走才是更稳的策略。如果你发现微调 epoch 数太少，精度还没稳住，可以把 5 个 epoch 加到 10 个，趋势不变。

3.4 实测数据：剪枝精度和微调收益

在固定随机种子和上述超参下，我这里跑出的典型数据如下：

看这个表第一反应应该是极端稀疏度下微调收益非常显著。95% 稀疏度意味着 44.4k 参数里只剩 2.2k 非零，剪完直接掉到 93.12%，微调一轮还能拉回 97.62%。这就是权重级剪枝的魅力：结构越自由，冗余挖掘得越干净。

如果换成逐层 topk，同样 90% 稀疏度，剪后精度大概率掉到 97% 以下，微调后也在 98.5% 附近徘徊，全局剪枝的优势在这个小模型上表现得很明显。

3.5 真正让模型“变小”：保存 mask 或稀疏索引

直接 torch.save(model.state_dict()) 的话，文件体积不会有任何变化。原因是 state_dict 里存的还是完整 shape 的张量，只是值变成 0，0 并没有在磁盘上省出空间。

要让剪枝成果传导到存储端，至少做两层改造：剪枝完成后保存 mask 和非零坐标，或者直接转成 CSR/CSC 稀疏格式；推理时再基于稀疏结构，配合支持稀疏运算的算子计算。给一个简化版保存思路：

import numpy as np

def save_sparse_state(model, masks, path): sparse_dict = {} for name, p in model.named_parameters(): if name in masks: indices = p.data.abs() > 0 vals = p.data[indices].cpu().numpy() idxs = indices.cpu().numpy().astype(np.int64) sparse_dict[name] = {"vals": vals, "idx": idxs} else: sparse_dict[name] = p.data.cpu().numpy() torch.save(sparse_dict, path)

这种保存方式在 90% 稀疏度下体积能压缩到原来的五分之一左右，95% 时接近十分之一。注意，这是“存储”层面的变小，和“推理速度”是两回事，下面会展开讲。

4. 高频踩坑与排查思路

4.1 一次性剪太猛，模型直接变成随机猜测

第一次跑 topk 剪枝的人，十有八九会先把 sparsity 设到 0.95 试一次，然后发现评估精度掉到两位数。这个现象太常见了，原因也简单：一次性剪掉 95% 参数，剩下的权重根本不足以表达原任务，微调也很难从随机猜测状态拉回来。

解决方案是迭代剪枝，也叫逐步剪枝。每次只增加 10% 的稀疏度，中间插入几轮微调，让剩余权重重新分布到更重要的位置，然后再继续。代码骨架如下：

target = 0.9 current = 0.0 for round_idx in range(10): current = min(target, current + 0.1) masks = global_topk_prune(model, current) apply_weight_mask(model, masks)

optimizer = torch.optim.SGD(model.parameters(), lr=0.002, momentum=0.9) for epoch in range(3): train_one_epoch(model, train_loader, optimizer, masks=masks) print(f"round {round_idx + 1}: sparsity={calc_sparsity(model):.2f}")

这就是彩票假设在工程上的落地：先小幅剪一点，重训把剩余参数重新分配到重要位置，再剪一点，循环往复。实测到 95% 以上稀疏度时，迭代法的精度比单次剪枝高出 2~3 个百分点。

4.2 微调几轮后稀疏度悄悄回弹

典型现象是剪完打印稀疏度是 0.9，微调几轮后一数变成 0.88 甚至 0.85。原因不外乎两个：忘记在 backward 后对梯度做 mask，或者优化器带了 weight_decay。

我见过太多人只做 weight mask，训练循环里没加 apply_grad_mask，于是每个 step 被剪位置的梯度又参与更新，权重悄悄长回来。另一个隐蔽来源是 AdamW 的 weight decay 直接叠加在梯度上，step 后权重立刻变成非零。两个问题统一解法就是上文提到的：backward 之后 mask 梯度，optimizer.step 之后再 mask 权重，双保险。

另外，模型里有 BN 的话，剪枝后统计量会漂移。必要时用训练集重新跑一遍 forward 更新 running_mean 和 running_var，或者直接多微调几个 epoch，让 BN 自己适应新的权重分布。

4.3 torch.topk 报错或剪完稀疏度不对

这一类问题往往出在几个小细节上。keep 数量必须保证至少为 1，否则 k=0 直接报错，所以代码里写成 max(int(...), 1)。如果模型参数里混入 NaN 或 inf，topk 的排序结果会乱，剪出来的稀疏度和预期完全对不上，排查时先检查参数统计信息。还有 reshape 和 view 的区别，原始张量内存不连续时 view 会报错，统一用 reshape 更安全。最后，如果你发现剪完稀疏度和目标差距大于 1%，检查一下全局剪枝阈值等值边界的问题，那属于正常误差，通常不需要处理。

下面是一个快速排错命令：

for name, p in model.named_parameters(): if torch.isnan(p).any() or torch.isinf(p).any(): print(f"参数异常: {name}")

4.4 模型文件没变小，推理也没加速

这两件事是 topk 剪枝最常被误解的地方。存储方面，如果保存的还是 dense state_dict，文件体积不变，因为 0 权重仍然占据存储位。要让体积变小，必须保存非零值加索引，或直接转稀疏格式。

推理方面的误解更普遍：PyTorch 默认的 dense 算子在计算矩阵乘法时并不知道“哪里是 0”，仍然按完整形状计算。剪完的稀疏权重跑在普通 GPU 上，速度不但不会提升，反而可能因为稀疏检查逻辑变慢。要让剪枝兑现推理加速，必须部署到支持稀疏张量运算的推理后端或硬件平台，由平台把计算下发到稀疏算子。

所以接到需求时，我第一件事永远是问清楚：目标到底是减小存储，还是提升线上推理吞吐。如果只是部署体积卡住，topk 剪枝非常合适；如果想省推理延迟，就要考虑通道剪枝、结构剪枝或稀疏算子适配，不能拿 topk 硬套。

4.5 剪完精度没掉，但也没提升

这个现象其实没必要焦虑。前期冗余高的时候，剪枝加微调后精度基本和基线持平，这是正常的。中高稀疏度下偶尔会出现比基线高零点几个百分点的情况，一般解释为剪枝起到了隐式正则化的作用，帮模型绕开了一些可能过拟合的特征组合。

别把这类小幅提升当成剪枝的常规收益。把它理解为“剪枝不亏”就好，真正的价值在压缩比，不在那零点几个点的精度波动。

4.6 常见问题速查表

整理成一张速查表，定位更快：

5. 往深了走：迭代调度、组合玩法与边界场景

5.1 迭代式剪枝与重训练调度

迭代剪枝本质上就是剪-训-剪的循环，但调度策略会直接影响最终精度。我常用的配置是每轮增加 10%~15% 稀疏度，训练 3~5 个 epoch。小模型可以少一点，大模型要多训几个 epoch，否则上一轮剪枝的影响还没收敛，下一轮又来了。

学习率调度要更保守。剪枝刚完成那次微调，学习率通常取基线的 0.1~0.2 倍；后续轮次再慢慢恢复，最高不超过基线学习率的一半。这个节奏比固定低学习率整个流程跑下来效果更稳，因为模型在新一轮剪枝之后需要先“止血”，再“恢复”。

还有一个小习惯：每轮迭代之后打印完整记录，包含当前稀疏度、剪后精度、微调后精度。有了这份记录，模型调参时会省很多事。

5.2 组合拳：剪枝 + 量化 + 蒸馏

单纯的 topk 剪枝压约 10 倍体积之后，精度基本能守住，但还想继续压就得和别人一起上。我通常的顺序是先剪枝，重训恢复精度，再做 8bit 量化，最后补一层蒸馏微调。

原因在于，剪枝后的稀疏模型浮点误差与 dense 模型接近，量化误差也相对可控；反过来先量化再剪枝，量化误差和剪枝扰动叠加在一起，排查困难。蒸馏时的 teacher 就用基线 dense 模型，student 是剪枝加量化后的模型，用 KL 散度把 logits 教回来。这套组合拳做完，90% 稀疏度配合 8bit 量化，模型总体积能压到原来的四十分之一左右。

5.3 哪些场景不适合 topk 剪枝

不是所有模型都适合这套方法。实时推理场景如果部署平台没有稀疏算子支持，剪枝之后拿不到加速收益，等于白干。小模型本身冗余低，压缩空间有限，剪了反而容易掉点。NLP 模型 embedding 占比很高，剪 embedding 对语义破坏明显，收益又不大，通常建议跳过 embedding，只剪 attention 和 FFN 里的权重。

模型参数量特别大的极稀疏场景也要慎重：全局 topk 需要对整个模型参数做一次展平排序，在十亿参数模型上内存开销不低。此时可以把参数分成几段，分段算阈值再做合并，或者用近似分位数方法替代完整 topk，工程处理手法很多。

如果让我只留一条实践经验给后来者，那就是不要指望一把梭。topk 剪枝应该被当成一个可迭代、可观察、可恢复的流程，而不是一条命令结束的魔法操作。我自己每次实验都会把基线精度、每轮稀疏度、剪后精度、微调后精度单独记下来，不是为了写报告，而是为了在下一次把稀疏度调高时，能一眼看出是哪一步出了问题。这个记录习惯帮我绕开了很多“精度莫名其妙掉了”的排查环节，也算是最朴素但最有用的一条经验。