模型剪枝（Model Pruning）

训练好的深度学习模型动辄上亿参数，直接搬到手机、嵌入式设备上，既占内存又算得慢。剪枝就是冲着这个问题来的，它把那些贡献微弱的权重参数直接置零，让参数矩阵变稀疏，存储和计算都能省下来，顺带还能压制过拟合。其实决策树里也有剪枝，作用是正则化防止过拟合，深度学习里的剪枝思路和它一脉相承。

接下来，我们讲清楚 PyTorch 里剪枝到底怎么落地，从随机剪枝、按范数剪枝到跨层的全局剪枝，再到怎么把剪枝结果固化成正常权重，最后澄清几个最容易踩的坑。读这篇之前，你只需要知道神经网络的权重是什么，会用 PyTorch 搭一个最简单的全连接层就行。

1. 剪枝概述

我们先把剪枝到底做了什么讲清楚，后面的接口才有意义。剪枝的本质动作只有一个，把一部分权重参数设成 0，让原本稠密的参数矩阵变得稀疏。那些值为 0 的位置在有效计算里不再出力，模型存的东西变少，算的东西也变少。被优先砍掉的，通常是那些绝对值很小、对最终输出影响微弱的连接，这等于强迫模型用更少的连接去表达同样的功能，本身就带一点正则化的味道。

剪枝分结构化和非结构化两种，差别直接决定能不能真的省算力。非结构化是逐个元素挑，把绝对值小的零散权重置零，张量的形状一点没变，只是里面多了很多 0，要靠稀疏计算库才能把这些 0 跳过。结构化是整行、整列、整个通道地砍，砍完之后张量真的变小变矮，不依赖特殊库也能直接省计算。

我们拿一个具体的小权重矩阵走一遍，剪枝前后的差别就一目了然。设权重 W 是一个 2 行 3 列的矩阵，一共 6 个参数，现在按绝对值把最小的 2 个置零，最小的两个是 1 和 2：

位置	原值	是否置零	置零后
(0,0)	1	是	0
(0,1)	2	是	0
(0,2)	3	否	3
(1,0)	4	否	4
(1,1)	5	否	5
(1,2)	6	否	6

剪完的有效权重就是 [[0,0,3],[4,5,6]]，6 个参数里 2 个变成 0，稀疏率等于 2 除以 6，约等于 33.3%。这里要注意，原始权重一个都没删，只是被一张掩码乘成了 0，这正是 PyTorch 剪枝的工作方式。

2. 剪枝的流程

知道剪枝做什么了，我们再看它在整个训练流程里处于哪一步。剪枝不是从零开始训，而是先把模型训练到一个不错的水平，再动手砍。典型的循环是这样，先正常训练拿到效果好的模型，然后剪枝，剪完精度必然往下掉一点，所以要先评估一下掉了多少。掉得太多就拿数据微调，让剩下的权重把丢掉的能力补回来，补完还可以再剪一轮、再微调，直到精度和压缩率都满意。

这里有个关键细节，再次微调的时候，被剪枝置零过的参数不会被更新，它们始终保持为 0。也就是说剪枝决定砍掉的连接是冻住的，微调只会在剩下的活参数里继续调，不会让它们偷偷长回来。

3. 随机剪枝

PyTorch 把剪枝封装得很细，我们先看最朴素的随机剪枝。它提供 random_structured 和 random_unstructured 两个接口，区别在砍的粒度。random_structured 必须指定 dim 参数，沿着某个维度整段地砍，属于结构化剪枝。random_unstructured 则是在所有元素里随机挑，不用指定维度，属于非结构化剪枝。到底砍多少由 amount 控制，传小数表示比例，传整数表示具体砍几个。

下面这段代码先搭一个带权重和偏置的极简网络，再分别演示两种随机剪枝，你重点看剪完之后出现的 orig 和 mask：

import torch
import torch.nn as nn
from torch.nn.utils.prune import random_structured
from torch.nn.utils.prune import random_unstructured
from torch.nn.utils.prune import l1_unstructured
from torch.nn.utils.prune import ln_structured
from torch.nn.utils.prune import global_unstructured
from torch.nn.utils.prune import L1Unstructured
from torch.nn.utils.prune import RandomUnstructured
from torch.nn.utils.prune import RandomStructured
from torch.nn.utils.prune import LnStructured


# 用于剪枝的简单网络
class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()
        self.weight = nn.Parameter(torch.tensor([[1., 2., 3.], [4., 5., 6.]]))
        self.bias = nn.Parameter(torch.tensor([10., 20., 30.]))


def test01():

    model = Model()
    print(model.state_dict()['weight'])
    print(model.state_dict()['bias'])

    print('-' * 30)

    # 随机在某个 dim 上稀疏一整段参数
    m = random_structured(model, name='weight', dim=1, amount=2)
    print(m.state_dict()['weight_orig'])
    print(m.state_dict()['weight_mask'])

    print('-' * 30)

    # 随机稀疏掉指定数量的参数为 0
    model = Model()
    m = random_unstructured(model, name='weight', amount=1)
    print(m.state_dict()['weight_orig'])
    print(m.state_dict()['weight_mask'])


if __name__ == '__main__':
    test01()

跑出来的结果是这样。第一段 random_structured 在 dim=1 上砍 2 列，掩码变成 [[0,0,1],[0,0,1]]，也就是第 0 列和第 1 列整列被清零，第 2 列保留。第二段 random_unstructured 只砍 1 个元素，掩码 [[1,1,1],[1,1,0]]，即位置 (1,2) 那个 6 被随机选中置零：

tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([10., 20., 30.])
------------------------------
tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([[0., 0., 1.],
        [0., 0., 1.]])
------------------------------
tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([[1., 1., 1.],
        [1., 1., 0.]])

这里要盯住一个变化。剪枝做完之后，原来叫 weight 的参数名不见了，取而代之的是 weight_orig 和 weight_mask 两个。真正前向用的 weight 等于这两个逐元素相乘，mask 为 0 的地方就是被砍掉的位置。另外这个函数是直接改在原模型上的，返回值主要方便你去看那两个字段。

4. 范数剪枝

随机砍没有任何依据，更好的做法是按权重大小砍，把不重要的优先砍掉。PyTorch 提供了基于 L1 范数的 l1_unstructured，以及可指定范数的 ln_structured。前者把张量里所有元素按绝对值排序，把最小的一批置零，不看形状和结构，所以是非结构化的。后者沿指定的 dim 把张量切成若干段，逐段算范数，把范数最小的整段砍掉，因为它照顾了形状，所以是结构化的。

我们还是用 W 等于 [[1,2,3],[4,5,6]] 这个矩阵。l1_unstructured 传 amount=2，绝对值最小的是 1 和 2，掩码变成 [[0,0,1],[1,1,1]]。ln_structured 传 n=1、dim=0，意思是按行算 L1 范数，第一行 1+2+3 等于 6，第二行 4+5+6 等于 15，砍掉范数小的第一行，掩码就是 [[0,0,0],[1,1,1]]：

def test02():

    model = Model()
    print(model.state_dict())
    # 把 L1 范数最低的那几个权值置为 0
    m = l1_unstructured(model, name='weight', amount=2)
    print(m.state_dict()['weight_orig'])
    print(m.state_dict()['weight_mask'])


def test03():
    model = Model()
    print(model.state_dict())
    # 通过 dim 指定维度，n 指定该维度上的范数，把最小的整段置 0
    m = ln_structured(model, name='weight', amount=1, n=1, dim=0)
    print(m.state_dict()['weight_orig'])
    print(m.state_dict()['weight_mask'])


if __name__ == '__main__':
    test02()
    print('-' * 30)
    test03()

OrderedDict([('weight', tensor([[1., 2., 3.],
        [4., 5., 6.]])), ('bias', tensor([10., 20., 30.]))])
tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([[0., 0., 1.],
        [1., 1., 1.]])
------------------------------
OrderedDict([('weight', tensor([[1., 2., 3.],
        [4., 5., 6.]])), ('bias', tensor([10., 20., 30.]))])
tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([[0., 0., 0.],
        [1., 1., 1.]])

5. 全局剪枝

前面都是盯着某一层的某一个权重单独剪，全局剪枝把视野放大到整个模型。它一次性指定要参与剪枝的所有参数，从这些参数的总量里统一挑出 amount 个最弱的置零，而不是每层各砍各的。pruning_method 传的是前面那几类剪枝方法，注意这里传的是类名而不是函数名，比如 RandomUnstructured、L1Unstructured 都可以。

def test04():

    model = Model()
    print(model.state_dict())

    # 定义要参与剪枝的参数，这里同时剪 bias 和 weight
    prune_param = ((model, 'bias'), (model, 'weight'))
    # pruning_method 传剪枝类名，从全局视角把 amount 个参数置 0
    global_unstructured(parameters=prune_param,
                        pruning_method=RandomUnstructured,
                        amount=2)

    print(model.state_dict())


if __name__ == '__main__':
    test04()

这次同时剪了 bias 和 weight 两组参数，一共 9 个数里随机挑 2 个置零。输出里 bias_mask 和 weight_mask 同时出现，加起来正好 2 个位置是 0，下面这次运行里 bias 的第一个数 10 和 weight 里的 4 被选中：

OrderedDict([('weight', tensor([[1., 2., 3.],
        [4., 5., 6.]])), ('bias', tensor([10., 20., 30.]))])
OrderedDict([('bias_orig', tensor([10., 20., 30.])), ('weight_orig', tensor([[1., 2., 3.],
        [4., 5., 6.]])), ('bias_mask', tensor([0., 1., 1.])), ('weight_mask', tensor([[1., 1., 1.],
        [0., 1., 1.]]))])

6. 固化结果

剪枝做完，模型里其实多出来一组 mask 参数，怎么把它变回一个干净的权重呢？PyTorch 剪枝本质上是给权重包了一层，weight 被拆成 weight_orig 和 weight_mask，前向时两者相乘才是真正的 weight。remove 函数就是把这层伪装拆掉，把 orig 乘 mask 固化成一个真正的 weight，从此不再有 mask 的概念。is_pruned 用来查一个模块当前到底有没有被剪过。

def test05():

    from torch.nn.utils.prune import remove
    import torch.nn.utils.prune as prune

    model = Model()
    print(model.state_dict(), prune.is_pruned(model))  # False
    m = l1_unstructured(model, name='weight', amount=2)
    print(model.state_dict(), prune.is_pruned(model))  # True
    # 把 name_orig 和 name_mask 合并回 name，去除剪枝伪装
    remove(model, name='weight')
    print(model.state_dict(), prune.is_pruned(model))  # False

7. 常见误区

第一个坑，以为把权重置零就自动加速了。PyTorch 的 prune 只是把权重乘上 mask 变成 0，稠密矩阵乘法照样把所有元素从头到尾算一遍，那些 0 并没有被跳过。真要省算力，还得把权重转成稀疏张量，或者自己写逻辑跳过零元素，光调剪枝函数不会让模型变快。

第二个坑，以为微调时被砍掉的参数会复活。mask 一直挡在那里，哪怕反向传播算回了梯度，这些位置也保持 0，剪枝的决定是冻住的。

第三个坑，amount 传小数还是整数搞混。传小数是比例，比如 0.3 表示砍掉三成，传整数是个数，比如 2 就砍两个，传错了砍的数量能差出一个数量级。

第四个坑，以为剪枝函数只返回结果不动原模型。它是原地修改的，传进去的那个 model 本身已经被改了，返回值主要让你方便查看 mask。

我们回到开头的问题。剪枝把不重要的权重置零换稀疏，结构化整段砍能直接缩小张量省计算，非结构化逐元素砍则要靠稀疏库才真的跳过零。整套流程是先训练出好模型，再剪枝，掉点精度就微调，被砍掉的参数全程冻住。PyTorch 里随机、按范数、全局三种剪枝覆盖了日常需求，最后用 remove 把 orig 和 mask 乘回成正常权重。最该记住的一条是，置零不等于自动加速，真正部署提速还得接稀疏计算。