4、非结构化剪枝实战：基于权重大小的剪枝、阈值设定、稀疏矩阵存储

好，咱们进入第四章。这一章，我打算带大家亲手做一次非结构化剪枝。

说实话，非结构化剪枝这名字听着挺唬人。但说白了，它的核心思想就一句话：把那些不重要的权重直接干掉。怎么判断不重要？看绝对值大小。绝对值小的，对模型输出贡献就小，砍掉它影响不大。

我在项目里遇到过不少次，模型太大，边缘设备跑不动。这时候，非结构化剪枝往往是第一板斧。它简单、粗暴，但效果立竿见影。

4.1 基于权重大小的剪枝：最朴素的思想

咱们先看最基础的做法。假设你训练好了一个模型，里面全是浮点数权重。现在，你设定一个阈值，比如 0.01。所有绝对值小于 0.01 的权重，统统置为 0。

就这么简单？对，就这么简单。但这里有个坑——阈值怎么定？

我个人的习惯是，先看一眼权重的分布。你可以把所有权重的绝对值拉出来，画个直方图。看看大部分权重集中在哪个区间。如果大部分权重都在 0.001 以下，那阈值设 0.01 可能就太狠了，会把模型剪废。

来看一段代码，这是 PyTorch 里的实现：

import torch
import torch.nn as nn

def magnitude_prune(model, threshold=0.01):
    """
    基于权重大小的剪枝
    """
    for name, param in model.named_parameters():
        if 'weight' in name:
            # 创建一个掩码，绝对值小于阈值的置为0
            mask = torch.abs(param.data) > threshold
            # 应用掩码
            param.data.mul_(mask.float())
    return model

# 使用示例
model = YourTrainedModel()
pruned_model = magnitude_prune(model, threshold=0.01)

你看，核心就两行：生成掩码，然后乘上去。掩码里，大于阈值的保留原值，小于阈值的变成 0。

注意： 这种剪枝是一次性的。剪完以后，模型精度通常会掉一截。你需要做微调（fine-tune）来恢复精度。我曾经试过不微调直接部署，结果精度掉了 5 个点，被老板骂了一顿。

4.2 阈值设定：不是拍脑袋，是有套路的

阈值怎么设？我见过有人拍脑袋设个 0.001，有人设 0.1。这都不对。

阈值设定的核心是：你要剪掉多少比例的权重？

比如，你想剪掉 50% 的权重。那你就把所有权重绝对值排个序，找到第 50 百分位的那个值，把它作为阈值。

这样做的好处是，剪枝比例是可控的。你可以在 30%、50%、70% 这几个比例上分别试一下，看看精度下降多少。然后选一个精度能接受的最高比例。

我给大家一个更实用的方法——基于百分位的阈值设定：

def percentile_threshold(model, sparsity=0.5):
    """
    根据目标稀疏度自动计算阈值
    sparsity: 目标稀疏度，0.5 表示剪掉50%的权重
    """
    all_weights = []
    for name, param in model.named_parameters():
        if 'weight' in name:
            all_weights.append(param.data.view(-1))
    
    # 合并所有权重
    all_weights = torch.cat(all_weights)
    # 计算阈值
    threshold = torch.quantile(torch.abs(all_weights), sparsity)
    return threshold

# 使用
target_sparsity = 0.6  # 剪掉60%
thresh = percentile_threshold(model, target_sparsity)
pruned_model = magnitude_prune(model, threshold=thresh)
小技巧： 不同层的权重分布可能差异很大。比如第一层卷积的权重通常比较大，而全连接层的权重可能很小。我建议逐层设定阈值，而不是全局统一一个阈值。这样能保留更多重要信息。

4.3 稀疏矩阵存储：剪完以后怎么办？

剪完了，模型里全是 0。但问题来了——这些 0 仍然占着内存。一个 float32 的 0，和 float32 的 0.5，占的内存一样大。那剪枝的意义在哪？

嗯，这里就要引入稀疏矩阵存储了。

说白了，我们只存非零元素的位置和值。常见的格式有两种：

格式 说明 适用场景
COO (Coordinate Format) 存三个数组：行索引、列索引、值 简单直观，适合小规模
CSR (Compressed Sparse Row) 存三个数组：值、列索引、行偏移 更节省内存，适合大规模

在 PyTorch 里，你可以直接用 torch.sparse 来操作。但要注意，不是所有算子都支持稀疏张量。我踩过这个坑——剪完以后想直接推理，结果报错说算子不支持稀疏输入。

所以，实际部署时，我们通常的做法是：

  1. 训练时：用密集张量，但权重被置为 0
  2. 推理时：转换成稀疏格式，或者直接用自定义的稀疏算子

来看一个转换示例：

import torch

def convert_to_sparse(tensor):
    """
    将密集张量转换为COO格式的稀疏张量
    """
    # 找出非零元素的索引
    indices = torch.nonzero(tensor, as_tuple=False).t()
    # 取出非零元素的值
    values = tensor[indices[0], indices[1]]
    # 创建稀疏张量
    sparse_tensor = torch.sparse_coo_tensor(
        indices, values, tensor.size()
    )
    return sparse_tensor

# 假设 weight 是剪枝后的权重
weight = torch.randn(4, 4)
weight[torch.abs(weight) < 0.5] = 0  # 模拟剪枝
sparse_weight = convert_to_sparse(weight)

print(f"密集张量占用: {weight.element_size() * weight.numel()} bytes")
print(f"稀疏张量占用: {sparse_weight.element_size() * sparse_weight._nnz() * 3} bytes")
# 注意：稀疏张量存储了索引+值，大约3倍于非零元素
关键点： 只有当稀疏度超过 70% 时，稀疏存储才有明显的内存优势。如果只剪了 20%，用稀疏存储反而更占内存。因为你要额外存索引信息。

4.4 实战中的避坑指南

最后，我分享几个实战中遇到的坑：

  • 剪枝后一定要微调：我曾经图省事，剪完直接部署。结果精度掉了 8 个点，用户反馈说模型变傻了。后来加了 3 个 epoch 的微调，精度基本恢复。
  • 注意 Batch Normalization 层：BN 层的参数不要剪。我见过有人把 BN 的 gamma 也剪了，结果模型输出全乱套。
  • 稀疏算子要提前验证：不是所有硬件都支持稀疏矩阵运算。我在某款边缘芯片上吃过亏——它不支持稀疏卷积，我只能把稀疏矩阵转回密集再算，结果内存没省下来，速度还慢了。

好了，这一章的内容就到这里。非结构化剪枝虽然简单，但它是后续所有剪枝技术的基础。下一章，我们会讲结构化剪枝，那才是真正能加速推理的硬核技术。

← 上一章 📖 返回目录 下一章 →