文章目录
- 一、剪枝基础
- 二、剪枝的分类
- 1.按照剪枝粒度分类( 粗 ➡️ 细 )
- 2.按照剪枝重要性准则分类
- a. 基于权重大小的剪枝(Weight Magnitude Pruning)
- b. 基于梯度的剪枝(Gradient-based Pruning)
- c. 基于敏感度的剪枝(Sensitivity-based Pruning)
- d. 基于熵的剪枝(Entropy-based Pruning)
- e. 基于L1/L2范数的剪枝(L1/L2 Norm-based Pruning)
- f. 基于频率的剪枝(Frequency-based Pruning)
- g. 基于贝叶斯方法的剪枝(Bayesian-based Pruning)
- h. 基于重要性评分的剪枝(Importance Score-based Pruning)
- 3.按照剪枝的策略分类
- 三、剪枝比率确定(延迟和精确性和模型大小间的权衡)
- 四、剪枝后的微调
- 五、思考总结
一、剪枝基础
1.剪枝定义
剪枝(Pruning)是指通过移除神经网络中的不必要部分(如权重、神经元或层)来减少模型的复杂度和运行成本。剪枝的目标是提高模型的计算效率、降低模型的存储需求以及减少能耗,同时尽量保持模型的原始性能。
2.剪枝时机
- 训练:首先,需要训练一个初始的大模型,通常是为了达到足够的性能水平。
- 剪枝:使用某种评估方法(如:权重的绝对值、梯度信息等)来确定模型中各个参数的重要性。根据评估结果,剪枝掉不重要的参数或连接,可以是结构化的或非结构化的。
- 微调:进行剪枝后,需要进行一定的修正和微调,以确保模型的性能不会显著下降
二、剪枝的分类
1.按照剪枝粒度分类( 粗 ➡️ 细 )
A.结构化剪枝
a. 隐藏层剪枝(Layer Pruning)
定义:移除神经网络中的整个层(如卷积层、全连接层等),以减少模型的深度和计算复杂度。
特点:
- 粒度:最粗粒度的剪枝方法。
- 优点:大幅度减少参数和计算量。
- 缺点:可能对模型性能影响较大,需要精细选择要移除的层。
import torch.nn as nn
class OriginalModel(nn.Module):
def __init__(self):
super(OriginalModel, self).__init__()
self.layer1 = nn.Linear(784, 256)
self.layer2 = nn.Linear(256, 128)
self.layer3 = nn.Linear(128, 10)
def forward(self, x):
x = torch.relu(self.layer1(x))
x = torch.relu(self.layer2(x))
x = self.layer3(x)
return x
class PrunedModel(nn.Module):
def __init__(self):
super(PrunedModel, self).__init__()
self.layer1 = nn.Linear(784, 256)
self.layer2 = nn.Linear(256, 10) # 移除了原来的layer3
def forward(self, x):
x = torch.relu(self.layer1(x))
x = self.layer2(x)
return x
# 使用PrunedModel代替OriginalModel进行后续训练和推理
应用场景:适用于冗余层较多的深度神经网络。
b. 通道剪枝(Channel Pruning)
定义:移除卷积神经网络中的整个卷积通道(也称特征图或滤波器),这相当于移除卷积层中的某些卷积核。
特点:
- 粒度:较粗粒度,移除的是网络中的整个卷积通道。
- 优点:显著减少计算量和内存占用,较好保持网络结构的连贯性。
- 缺点:需要准确评估每个通道的重要性。
import torch
import torch.nn as nn
import torch.optim as optim
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.fc1 = nn.Linear(32*8*8, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.relu(self.conv2(x))
x = torch.flatten(x, 1)
x = self.fc1(x)
return x
# 示例:移除conv2中的一半通道
def prune_channels(model, layer, num_channels):
_, indices = torch.topk(layer.weight.data.norm(2, dim=(1, 2, 3)), num_channels, largest=False)
layer.weight.data = torch.index_select(layer.weight.data, 0, indices)
layer.bias.data = torch.index_select(layer.bias.data, 0, indices)
# 创建模型实例
model = SimpleCNN()
# 剪枝conv2层,移除一半的通道
prune_channels(model, model.conv2, model.conv2.out_channels // 2)
应用场景:适用于卷积神经网络(CNN)中,有助于减少计算复杂度并加速推理。
c. 卷积核剪枝(Kernel Pruning)
定义:移除卷积操作中的整个卷积核,即移除卷积层中的某个卷积滤波器。
特点:
- 粒度:中等粒度,移除的是卷积层中的单个卷积核。
- 优点:减少卷积操作的计算量,保持卷积层的基本结构。
- 缺点:需要对卷积核的重要性进行评估和选择。
def prune_kernels(model, layer, num_kernels):
_, indices = torch.topk(layer.weight.data.norm(2, dim=(1, 2, 3)), num_kernels, largest=False)
layer.weight.data = torch.index_select(layer.weight.data, 0, indices)
layer.bias.data = torch.index_select(layer.bias.data, 0, indices)
# 剪枝conv1层,移除一半的卷积核
prune_kernels(model, model.conv1, model.conv1.out_channels // 2)
应用场景:适用于需要细化至卷积核级别的剪枝需求,以减少计算复杂度和参数数量。
B.非结构化剪枝
a. 核内权重剪枝(Intra Kernel Weight Pruning)
定义:移除卷积核内部的部分权重,而不是整个卷积核。
特点:
- 粒度:细粒度,针对卷积核内部的个别权重进行剪枝。
- 优点:精细控制剪枝粒度,减少计算量的同时尽量保留模型性能。
- 缺点:复杂度较高,需要对卷积核内的每个权重进行评估。
def prune_intra_kernel_weights(model, layer, threshold):
mask = layer.weight.data.abs() > threshold
layer.weight.data *= mask.float()
# 剪枝conv1层,移除小于某个阈值的权重
prune_intra_kernel_weights(model, model.conv1, 0.1)
应用场景:适用于需要非常精细的剪枝操作,以在最大程度减少计算量的同时保持模型性能。
b. 单个权重剪枝(Weight Pruning)
定义:移除神经网络中的单个权重,这通常是基于权重的绝对值大小进行剪枝。
特点:
- 粒度:最细粒度,移除单个权重。
- 优点:可以最大化减少参数数量,保持模型性能。
- 缺点:计算复杂度高,需要对大量权重进行评估和选择。
import torch.nn.utils.prune as prune
def prune_weights(model, layer, amount):
prune.l1_unstructured(layer, name='weight', amount=amount)
# 剪枝conv1层,移除10%的权重
prune_weights(model, model.conv1, 0.1)
应用场景:适用于需要最大化减少模型参数数量的场景,同时希望尽量保持原有模型的性能。
C.特殊情况
a. Token剪枝(Token Pruning)
定义:在自然语言处理任务中,移除输入序列中的不重要token,以减少输入序列长度和计算复杂度。
特点:
- 粒度:针对输入序列中的单个token。
- 优点:显著减少输入序列长度和计算复杂度。
- 缺点:需要仔细评估每个token的重要性,确保对模型性能的影响最小。
import torch
import torch.nn as nn
class SimpleTransformer(nn.Module):
def __init__(self):
super(SimpleTransformer, self).__init__()
self.embedding = nn.Embedding(1000, 256)
self.transformer = nn.Transformer(d_model=256, nhead=8, num_encoder_layers=2)
self.fc = nn.Linear(256, 10)
def forward(self, x):
x = self.embedding(x)
x = self.transformer(x)
x = torch.mean(x, dim=0)
x = self.fc(x)
return x
# 示例:移除不重要的token
def prune_tokens(tokens, importance_scores, threshold):
return tokens[importance_scores > threshold]
# 创建模型实例
model = SimpleTransformer()
# 假设我们有一些token及其重要性评分
tokens = torch.tensor([1, 2, 3, 4, 5])
importance_scores = torch.tensor([0.1, 0.4, 0.2, 0.8, 0.3])
# 剪枝token,移除重要性低于阈值的token
pruned_tokens = prune_tokens(tokens, importance_scores, 0.3)
应用场景:适用于基于Transformer的NLP模型,特别是在需要加速推理的场景中。
2.按照剪枝重要性准则分类
根据权重或神经元对模型输出的重要性进行剪枝。重要性可以通过各种方法计算,如梯度、Hessian 矩阵等。
a. 基于权重大小的剪枝(Weight Magnitude Pruning)
- 定义:根据权重的绝对值大小来判断权重的重要性,通常剪去绝对值较小的权重。
例如:基于神经元连接权重数值幅度的方法 - 优点:计算简单,容易实现。
- 缺点:可能会忽略一些对模型性能影响较大的小权重。
import torch
import torch.nn as nn
def weight_magnitude_prune(layer, threshold):
mask = layer.weight.data.abs() > threshold
layer.weight.data *= mask.float()
# 创建一个示例模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 实例化模型
model = SimpleModel()
# 对fc1层进行基于权重大小的剪枝
weight_magnitude_prune(model.fc1, threshold=0.1)
b. 基于梯度的剪枝(Gradient-based Pruning)
- 定义:通过分析权重的梯度来判断其重要性,即利用训练过程中权重的变化率来进行剪枝。
例如:使用目标函数对参数求二阶导数表示参数的贡献度的方法 - 优点:能够更好地反映权重在损失函数中的作用。
- 缺点:需要额外计算梯度信息,计算复杂度较高。
import torch
import torch.nn as nn
import torch.optim as optim
def gradient_based_prune(model, data, target, layer, threshold):
model.eval()
optimizer = optim.SGD(model.parameters(), lr=0.01)
output = model(data)
loss = nn.CrossEntropyLoss()(output, target)
optimizer.zero_grad()
loss.backward()
mask = layer.weight.grad.abs() > threshold
layer.weight.data *= mask.float()
# 假设我们有一些输入数据和目标标签
data = torch.randn(1, 784)
target = torch.tensor([3])
# 对fc1层进行基于梯度的剪枝
gradient_based_prune(model, data, target, model.fc1, threshold=0.01)
c. 基于敏感度的剪枝(Sensitivity-based Pruning)
-
定义:通过评估某个权重或神经元对整体模型性能(如精确度)的影响来进行剪枝,重要性低的部分被剪去。
-
优点:能够确保剪枝后性能影响较小。
-
缺点:需要进行多次评估,计算成本较高。
import torch
import torch.nn as nn
def sensitivity_based_prune(model, data, target, layer, threshold):
model.eval()
output = model(data)
loss = nn.CrossEntropyLoss()(output, target)
loss.backward()
sensitivity = layer.weight.grad.abs().max(dim=1).values
mask = sensitivity > threshold
layer.weight.data *= mask.float().unsqueeze(1)
# 对fc1层进行基于敏感度的剪枝
sensitivity_based_prune(model, data, target, model.fc1, threshold=0.01)
d. 基于熵的剪枝(Entropy-based Pruning)
-
定义:利用熵测量来判断权重在信息传递过程中的重要性,熵值低的部分被剪去。
-
优点:能够从信息论角度评估重要性。
-
缺点:计算复杂度较高,可能需要额外的理论背景知识。
import torch
import torch.nn as nn
def entropy_based_prune(layer, threshold):
entropy = -layer.weight.data * torch.log(torch.abs(layer.weight.data) + 1e-10)
mask = entropy.abs() > threshold
layer.weight.data *= mask.float()
# 对fc1层进行基于熵的剪枝
entropy_based_prune(model.fc1, threshold=0.01)
e. 基于L1/L2范数的剪枝(L1/L2 Norm-based Pruning)
-
定义:通过计算权重的L1或L2范数来评估其重要性,通常剪去范数较小的权重或神经元。
-
优点:实现简单,数学计算清晰。
-
缺点:仅依据权重的大小,可能忽略权重间的相互关系。
import torch
import torch.nn as nn
def l1_norm_prune(layer, threshold):
norm = layer.weight.data.abs().sum(dim=1)
mask = norm > threshold
layer.weight.data *= mask.float().unsqueeze(1)
# 对fc1层进行基于L1范数的剪枝
l1_norm_prune(model.fc1, threshold=0.01)
def l2_norm_prune(layer, threshold):
norm = layer.weight.data.pow(2).sum(dim=1).sqrt()
mask = norm > threshold
layer.weight.data *= mask.float().unsqueeze(1)
# 对fc1层进行基于L2范数的剪枝
l2_norm_prune(model.fc1, threshold=0.01)
f. 基于频率的剪枝(Frequency-based Pruning)
-
定义:分析某个权重或神经元在训练过程中被激活的频率,频率低的部分被剪去。
-
优点:能够识别训练过程中不常用的权重。
-
缺点:需要对训练过程进行详细记录和分析。
import torch
import torch.nn as nn
def frequency_based_prune(layer, activation_counts, threshold):
mask = activation_counts > threshold
layer.weight.data *= mask.float().unsqueeze(1)
# 示例:假设我们有一些激活计数
activation_counts = torch.tensor([10, 20, 5, 15, 50, 5, 30, 10, 25, 0])
# 对fc1层进行基于频率的剪枝
frequency_based_prune(model.fc1, activation_counts, threshold=10)
g. 基于贝叶斯方法的剪枝(Bayesian-based Pruning)
-
定义:通过贝叶斯统计方法评估权重的重要性,利用概率分布来进行剪枝。
-
优点:提供了统计学上的稳健性。
-
缺点:计算复杂度高,理论要求高。
import torch
import torch.nn as nn
import torch.distributions as dist
def bayesian_based_prune(layer, threshold):
posterior = dist.Normal(layer.weight.data.mean(), layer.weight.data.std())
mask = posterior.log_prob(layer.weight.data) > threshold
layer.weight.data *= mask.float()
# 对fc1层进行基于贝叶斯方法的剪枝
bayesian_based_prune(model.fc1, threshold=-3.0)
h. 基于重要性评分的剪枝(Importance Score-based Pruning)
-
定义:为每个权重或神经元分配一个重要性评分,根据评分进行剪枝。
-
优点:灵活性高,可以结合多种因素。
-
缺点:评分标准的选择和计算可能较为复杂。
import torch
import torch.nn as nn
def importance_score_prune(layer, scores, threshold):
mask = scores > threshold
layer.weight.data *= mask.float().unsqueeze(1)
# 假设我们有一些重要性评分
importance_scores = torch.tensor([0.1, 0.4, 0.3, 0.7, 0.2, 0.5, 0.6, 0.8, 0.9, 0.05])
# 对fc1层进行基于重要性评分的剪枝
importance_score_prune(model.fc1, importance_scores, threshold=0.3)
3.按照剪枝的策略分类
a.一次性剪枝(One-shot Pruning)
定义:在训练前或训练中期进行一次性剪枝。
优点:实现简单,节省训练时间。
缺点:可能导致模型性能显著下降,需要后续的再训练。
b.逐步剪枝(Iterative Pruning)
定义:在训练过程中逐步进行剪枝,每次剪枝后继续训练。
优点:剪枝过程更温和,对模型性能影响较小。
缺点:需要较长的训练时间和多次迭代。
三、剪枝比率确定(延迟和精确性和模型大小间的权衡)
1.手动确定:敏感性分析
分析每一层的敏感性,移除敏感度较低的层,做法是采用某一准确度给各层设计阈值。
(看似L1是最高效的,但是L1的10%一定大于L2的1%吗?还得看各个模型层的大小);
还好目前的transformer模型每层具有很大的同质性,但卷积神经网络层差异就比较大。
当然这是假设以上层都是独立的,我们还得考虑层之间的相互性。
剪枝比率和加速比率并不是单纯的线性关系,比如100个通道,剪枝25%为何获得70%的加速:因为输入通道(0.75)*输出通道(0.75)。
2.自动化确定剪枝的办法
a.使用AutoML的AMC
b.基于规则的迭代和渐进方法选择每层的剪枝程度
c.有旁路依赖关系的模型怎么剪枝
四、剪枝后的微调
1.剪枝模型微调怎么调整学习率
单步剪枝(训练到收敛然后剪枝)和迭代剪枝(如训练20%剪枝再训练再剪枝)
逐步降低学习率:
剪枝后的模型可能需要重新适应新的参数空间,因此逐步降低学习率是一种常用的策略。这可以防止模型在微调过程中发生剧烈的参数变化,保证训练的稳定性。
使用学习率调度器:
学习率调度器可以根据预定义的规则自动调整学习率,如每隔几个epoch降低一次学习率,或者在损失不再降低时调整学习率。
热身阶段:
先使用较低的学习率进行一段时间的“热身”训练,然后再逐步提高学习率。
实验性调整:
通过实验找到一个合适的初始学习率。一般可以测试几个不同的学习率,选择性能最佳的那个。
五、思考总结
非结构化剪枝与硬件支持
- 非结构化剪枝:剪去单个权重,导致稀疏矩阵。
- 硬件支持:在GPU上效果不好,需要专用加速器(如EIE)。
不同网络层和架构的剪枝难度
- 高层冗余性更大:高层网络容易剪枝。
- 卷积层 vs 全连接层:卷积层冗余性少,剪枝难度大;全连接层即使大幅剪枝,精度仍能保持不变。
- ResNet、MobileNet vs VGG、AlexNet:前者剪枝难度大,后者剪枝容易。
神经元剪枝与权重剪枝对模型精度的影响
- 神经元剪枝:剪去整个神经元,容易损失精度。
- 权重剪枝:剪去个别权重,精度损失较小。
梯度中的稀疏度
- 训练阶段梯度稀疏:梯度中有大量稀疏性,可以利用剪枝。
剪枝时机之最新研究
- 传统流程:预训练 -> 剪枝 -> 微调。
- 最新研究:随机初始化 -> 剪枝 -> 训练,可能获得更高稀疏度和精度。