理论依据#xff1a;有研究发现#xff0c;在许多深度神经网络中#xff0c;大部分参数是接近于0的#xff0c;这些参数对模型最终的…1.引言

模型剪枝本质上是一种利用稀疏性来减少模型大小和计算量从而提高训练和推理效率的技术。

它为何会有效呢

理论依据有研究发现在许多深度神经网络中大部分参数是接近于0的这些参数对模型最终的性能贡献较小。

这也就意味着识别并移除那些对模型性能影响较小的参数可以减少模型的复杂度和计算成本并且不会影响到模型的准确性。

根据粒度不同剪枝可以有不同的方法并且不同的模型由于内部结构不同剪枝方法也有所区别。

就像卷积神经网络中可以对卷积核剪枝和通道进行剪枝而transformer模型中则可以针对不活跃的自注意力头进行剪枝。

权重级剪枝基于模式剪枝向量级剪枝

本文将以一个二维矩阵为例来分别介绍这三种基本粒度的剪枝为了方便观察我们先来讨论下矩阵的可视化。

torch

注为了后面剪枝的可视化需要我们会在显示时将0值元素与其它元素区分开。

def

将矩阵转换为0和非0两个类别并设置两个类别的颜色映射为tab20cax.imshow(tensor.cpu().numpy()

vmin0,

cmaptab20c)ax.set_title(title)ax.set_yticklabels([])ax.set_xticklabels([])#

cols

colork)plt.show()plot_tensor(weight,

weight)3.权重级剪枝

权重级剪枝又称为细粒度剪枝以单个权重为剪枝单位在具体操作时一般会定义一个规则来决定移除哪些值。

在下面这个方法中会通过目标张量与掩码相乘的方式将小于threshold阀值的权重都置为0。

def

tensor.mul(mask)注1torch.ge函数的作用是对张量中的每个权重值与给定阀值threshold进行逐元素比较大于阀值的会置为True反之则置为False函数运算结果是一个0False)和1(True)组成的掩码。

注2mul

用于对张量进行按元素乘法要求两个矩阵的形状完全相同请与矩阵乘法运算符区分开。

pruned_weight

pruned_weight)权重级剪枝不关心权重在网络中的位置灵活度最高可实现高压缩比。

但是它破坏了原有模型的结构现有硬件架构的计算方式通常无法对它进行加速需要特殊的硬件或软件才能利用剪枝后模型的稀疏性所以在目前通用的硬件上运行时速度并不能得到提升。

4.基于模式剪枝

基于模式剪枝通常是基于非常规则的N:M稀疏性进行剪枝它要求在M个连续权重中固定有N个非零值而其余元素均置为0。

下面我们将以2:4稀疏性为例子一步一步说明如何实现N:M稀疏性。

4.1

0.])用permutations函数生成patterns列表的所有可能排列并用set去重。

from

set(permutations(sequence.tolist()))

list(patterns)[(0.0,

1.0)]这个patterns中包含了长度为4恰好有两个1的所有可能排列模式。

def

torch.Tensor(list(set(permutations(patterns.tolist()))))return

patternstensor([[0.,

计算掩码的目的是为了找到每个权重分组每M个连续权重为一组的最佳稀疏模式。

(-1,

torch.IntTensor(tensor.shape).fill_(1).view(-1,

masktensor([[1,

dtypetorch.int32)重塑权重矩阵使其形状与掩码一致。

mat

0.8882]])对于每一行权重我们计算其绝对值与所有稀疏模式的点积得到每一行权重在每一种稀疏模式下的加权和。

mat_patterns

patterns.t())

1.4296]])我们只需要保留加权和最大的模式即可为每一行权重选择加权和最大的模式索引

pmax

dtypetorch.int32)最后将掩码视图重塑回原始权重的形状。

mask

dtypetorch.int32)同样将上面计算掩码的过程封装为一个函数。

def

torch.IntTensor(tensor.shape).fill_(1).view(-1,m)

tensor.shape[1]

torch.FloatTensor(tensor.shape[0],

tensor.shape[1]

torch.argmax(torch.matmul(mat.abs(),

patterns.t()),

16行N列每一行的点积操作都得到N种可能取点积最大值的元素下标mask[:]

patterns[pmax[:]]

pruned_pattern)可以看到2:4稀疏性将一半权重都置为了0每4个权重中保留了两个非0权重我们成功地在权重矩阵上应用了

2:4

虽然都属于非结构化索引但与前面基于权重的索引不同的是2:4结构的稀疏性可以被英伟达的稀疏张量核心加速。

在运算时稀疏矩阵W首先会被压缩压缩后的矩阵存储着非零的数据值而metadata则存储着对应非零元素在原矩阵W中的索引信息。

具体来说metadata会将W中非零元素的行号和列号压缩成两个独立的一维数组这两个数组就是metadata中存储的索引信息。

注N:M

稀疏性是一种有效的稀疏技术这种稀疏模式不仅减少了模型大小还提高了计算效率特别适用于硬件加速器如

GPU

tensor.clone()prune_weight[rows,

0prune_weight[:,

pruned_vector)经过上面的向量级剪枝后可以直接去掉一行一列整个权重矩阵的形状可以直接由[8,8]变为[7,7],

通常在进行向量级别的剪枝时需要对模型的所有层统一进行剪枝其目的是在整个模型中保持一致的稀疏结构以确保上下游各层中结构和计算的一致性。

因此这种方法被称为“全局剪枝”或“统一剪枝”。

小结本文主要介绍了权重级剪枝、基于模式剪枝和向量级剪枝三种不同粒度的剪枝方法并结合可视化的方式一步一步详细演示了每种剪枝方法的运算过程和剪枝前后权重矩阵的变化。

在实际场景中基于模式的剪枝越来越受到青睐因为其规则的稀疏模式可以充分利用硬件加速器的计算能力从而显著提高计算效率。

参考阅读