4、非结构化剪枝：权重剪枝、阈值设定、稀疏训练与重训练策略，以及PyTorch实现示例

非结构化剪枝（Unstructured Pruning）是模型压缩中最细粒度的一类方法，它直接作用于神经网络中的单个权重参数，通过将值较小的权重置为零，使权重矩阵变为稀疏矩阵。与结构化剪枝（如通道剪枝、滤波器剪枝）不同，非结构化剪枝不改变张量的形状和维度，因此不需要特殊的硬件或推理库即可实现理论上的压缩效果。然而，其稀疏性通常是随机的、不规则的，导致实际加速依赖于对稀疏矩阵运算有良好支持的底层库（如PyTorch的稀疏张量、DeepSparse、TensorRT的稀疏推理模式）。本章将系统讲解权重剪枝的核心原理、阈值设定方法、稀疏训练与重训练策略，并给出完整的PyTorch实现示例。

上图展示了非结构化剪枝从预训练模型到稀疏模型部署的完整链路。核心决策点在于阈值设定策略（全局、分层或按比例）以及剪枝后的重训练方案。下面我们将逐一展开每个环节的细节。

4.1 核心原理：从稠密到稀疏的数学本质

神经网络中的权重矩阵通常包含大量冗余参数。研究表明，现代深度网络（如ResNet、BERT）中超过90%的权重值接近于零，它们对最终预测的贡献微乎其微。非结构化剪枝的基本思想是：将权重矩阵中低于某个阈值的元素置为零，从而将稠密矩阵转化为稀疏矩阵。数学上，给定权重矩阵 \(W \in \mathbb{R}^{m \times n}\)，剪枝操作可表示为：

W_sparse = W * M
其中 M ∈ {0,1}^{m×n}，M_ij = 1 当且仅当 |W_ij| ≥ τ
τ 为剪枝阈值

这里的掩码矩阵 \(M\) 是二值的，它决定了哪些权重被保留、哪些被丢弃。剪枝后的模型在推理时，可以跳过与零权重相关的乘法运算，从而减少计算量。但需要注意的是，非结构化剪枝的稀疏模式是随机的，这意味着内存访问模式不规则，对硬件缓存不友好。因此，实际加速效果取决于推理框架是否支持稀疏矩阵运算（如CSR/CSC格式）以及硬件是否具备稀疏计算单元。

4.2 阈值设定：全局、分层与敏感度分析

阈值设定是非结构化剪枝中最关键的步骤，它直接决定了剪枝率（即被置零的权重比例）和模型精度之间的平衡。常用的阈值设定方法有以下三种：

在实际操作中，分层阈值通常优于全局阈值，因为深度网络的不同层对剪枝的敏感度差异显著。例如，第一层卷积核通常提取低级特征（边缘、纹理），对剪枝非常敏感；而最后的全连接层往往冗余度较高，可以承受更高的剪枝率。敏感度分析的具体做法是：固定其他层不剪，只对某一层施加不同剪枝率（如10%、30%、50%、70%、90%），观察验证集精度的下降幅度。通过这种方式，可以为每一层找到“精度开始急剧下降”的临界剪枝率。

4.3 稀疏训练与重训练策略

剪枝后的模型直接推理通常会导致精度显著下降，因此必须通过重训练（Fine-tuning）来恢复精度。根据训练方式的不同，可以分为以下两种策略：

4.3.1 剪枝后重训练（Post-training Fine-tuning）

这是最经典的方法：先对预训练模型进行剪枝，然后在原始数据集上以较小的学习率（如原学习率的1/10）继续训练若干轮。重训练时，必须保持掩码矩阵 \(M\) 固定不变，即被剪掉的权重在反向传播中梯度为零，不再更新。PyTorch中可以通过注册掩码为buffer或手动在每次梯度更新后重新应用掩码来实现。

# 重训练时的关键代码模式
for batch in dataloader:
    optimizer.zero_grad()
    output = model(batch)
    loss = criterion(output, target)
    loss.backward()
    # 关键：将掩码为0的位置的梯度置零
    for name, param in model.named_parameters():
        if name in mask_dict:
            param.grad.data.mul_(mask_dict[name])
    optimizer.step()

4.3.2 稀疏训练（Sparse Training）

稀疏训练是指从随机初始化或预训练权重开始，在训练过程中动态地更新掩码。代表性方法包括SR-STE（Sparse Refinement with Straight-Through Estimator）和RigL（Rigged Lottery）。这些方法每隔若干步会重新计算掩码：移除当前值最小的部分权重，同时恢复一部分被剪掉但梯度较大的权重。这种动态稀疏训练（Dynamic Sparse Training）能够在训练过程中自动发现最优的稀疏结构，通常比“先剪后训”获得更好的精度-压缩比平衡。

4.4 PyTorch实现示例：权重剪枝与重训练完整流程

下面给出一个完整的PyTorch实现示例，演示如何对ResNet-18进行非结构化剪枝并重训练。该示例包含：阈值计算、掩码生成、剪枝应用、重训练循环。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, datasets, transforms

# 1. 加载预训练模型
model = models.resnet18(pretrained=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 2. 定义剪枝函数（分层阈值）
def prune_model(model, prune_ratio=0.5):
    mask_dict = {}
    for name, param in model.named_parameters():
        if param.dim() > 1:  # 只剪权重，不剪bias
            # 计算该层权重的值
            abs_weight = torch.abs(param.data)
            # 计算阈值：按百分位数
            threshold = torch.quantile(abs_weight.flatten(), prune_ratio)
            # 生成掩码
            mask = (abs_weight > threshold).float()
            mask_dict[name] = mask
            # 应用掩码
            param.data.mul_(mask)
    return mask_dict

# 3. 执行剪枝
mask_dict = prune_model(model, prune_ratio=0.5)

# 4. 重训练配置
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# 5. 重训练循环（关键：梯度掩码）
def train_one_epoch(model, dataloader, optimizer, criterion, mask_dict):
    model.train()
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        # 应用掩码：将剪掉权重的梯度置零
        for name, param in model.named_parameters():
            if name in mask_dict:
                param.grad.data.mul_(mask_dict[name])
        optimizer.step()

# 6. 训练循环（假设已有dataloader）
for epoch in range(5):
    train_one_epoch(model, train_loader, optimizer, criterion, mask_dict)
    scheduler.step()
    # 验证精度...

上述代码中，prune_model函数使用torch.quantile计算分层阈值，将值低于该阈值的权重置零。重训练时，通过param.grad.data.mul_(mask_dict[name])确保被剪掉的权重梯度为零，从而在后续优化步骤中不会被更新。这是实现非结构化剪枝重训练的核心技巧。

4.5 案例拆解：ResNet-18在CIFAR-10上的剪枝实战

我们以一个具体的实验来展示非结构化剪枝的效果。实验设置如下：

• 模型：ResNet-18，预训练于ImageNet，微调至CIFAR-10（输入尺寸调整为32×32）。
• 数据集：CIFAR-10，50,000张训练图，10,000张测试图。
• 剪枝配置：分层阈值，全局剪枝率50%（即每层剪掉50%的权重）。
• 重训练：学习率0.001，SGD优化器，训练5个epoch。

从结果可以看出：剪枝后直接推理精度下降了13.8个百分点，而经过5个epoch的重训练后，精度恢复到90.8%，仅比原始模型低1.6个百分点。这说明重训练对于恢复精度至关重要。进一步分析发现，第一层卷积的剪枝率虽然也是50%，但其对精度的影响远大于最后一层。如果采用敏感度分析，将第一层剪枝率降至10%，最后一层提升至80%，整体精度可以恢复到91.5%以上。

4.6 常见误区与总结

在实际项目中，工程师经常陷入以下误区：

1. 误区一：认为剪枝后模型大小一定减小。实际上，如果不使用稀疏存储格式（如CSR），模型文件大小不会改变，因为零权重仍然占用存储空间。只有配合稀疏编码或专用推理库才能实现存储和计算的双重压缩。
2. 误区二：重训练时忘记固定掩码。如果掩码不固定，被剪掉的权重会在梯度更新中重新变为非零值，导致稀疏度逐渐降低，最终失去剪枝效果。
3. 误区三：对所有层使用相同的剪枝率。这种做法往往导致敏感层（如第一层、最后的全连接层）精度崩溃，而冗余层（如中间卷积层）剪枝不足。
4. 误区四：忽略推理框架的稀疏支持。如果部署环境不支持稀疏矩阵运算，非结构化剪枝不仅无法加速，反而可能因为不规则的内存访问而变慢。

总结而言，非结构化剪枝是一种理论优雅、实现灵活的参数压缩技术。它的核心在于阈值设定与重训练策略的配合。对于大多数场景，推荐采用“分层阈值 + 敏感度分析 + 固定掩码重训练”的组合方案。在PyTorch中，实现非结构化剪枝的代码量不大，但需要仔细处理掩码的生成、应用和梯度屏蔽。当项目对推理延迟有硬性要求且硬件支持稀疏计算时，非结构化剪枝是值得优先考虑的技术路线。