从零学习大模型（十二）-----基于梯度的重要性剪枝（Gradient-based Pruning）

最新推荐文章于 2025-09-08 18:15:08 发布

原创最新推荐文章于 2025-09-08 18:15:08 发布 · 2.1k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#学习 #剪枝 #算法 #深度学习 #transformer #chatgpt #人工智能

大语言模型专栏收录该内容

19 篇文章

订阅专栏

部署运行你感兴趣的模型镜像

梯度的重要性定义

权重重要性（Weight Importance）

权重重要性通常指的是某个权重参数在模型输出中的影响程度。权重重要性的评估通常基于以下几个方面：

绝对值：一个常见的方法是直接使用权重的绝对值作为其重要性指标。权重越大，表示其对模型输出的影响越大，因此可以认为其重要性越高。
对模型性能的影响：通过在剪枝前后比较模型的性能，可以间接评估某个权重的重要性。若剪除某个权重后，模型性能显著下降，则说明该权重是重要的。

梯度的重要性（Gradient Importance）

梯度的重要性则是基于模型的训练过程来评估权重的重要性，具体体现在以下几个方面：

梯度的计算：在反向传播过程中，每个权重的梯度值反映了该权重对损失函数的影响。如果一个权重的梯度较大，意味着它在当前训练样本下对模型的输出变化有较大贡献。
动态性：与静态的权重重要性不同，梯度的重要性是动态的，会随着训练过程中的样本和损失函数的变化而变化。这意味着在不同的训练阶段或针对不同的数据样本，某些权重的梯度可能会显示出不同的“重要性”。

区别与联系

静态 vs. 动态：权重重要性通常是静态评估，基于权重本身的值；而梯度的重要性是动态的，依赖于当前的训练状态和样本。
应用场景：在剪枝时，可以综合考虑两者。某些方法可能首先基于权重重要性进行初步剪枝，然后使用梯度重要性进行精细调整，以优化剪枝后的模型性能。
互补性：两者可以结合使用，提供更全面的剪枝策略。例如，优先剪除权重绝对值小且梯度小的权重，以最小化对模型性能的影响。

基于梯度的重要性剪枝的工作原理

1.梯度的重要性评估

**损失函数：**假设有一个损失函数 $L$ ，通常是模型输出与真实标签之间的差异，比如均方误差或交叉熵：
$L=1N∑i=1Nl(yi,y^i) L = \frac{1}{N} \sum_{i=1}^{N} l(y_i, \hat{y}_i)$
其中 $y_i$ 是真实标签， $y^i\hat{y}_i$ 是模型预测。

**权重的梯度计算：**在反向传播过程中，计算每个权重 $w_j$ 的梯度：
$g_j = \frac{\partial L}{\partial w_j}$
这里 $g_j$ 是与权重 $w_j$ 相关的梯度。

2.确定剪枝阈值

**重要性度量：**通常用梯度的绝对值作为权重的重要性指标：
$I_j = |g_j|$
其中 $I_j$ 是权重 $w_j$ 的重要性。

**剪枝阈值：**设定一个阈值 $θ\theta$ ，可以选择某个比例的权重进行剪枝。例如，保留最重要的 $k%k\%$ 的权重：
$\text{Threshold} = \text{quantile}(I, 1 - k)$

3.执行剪枝

**剪除不重要的权重：**设定剪枝规则，所有重要性低于阈值的权重被剪除：
$w_j' = \begin{cases} 0 & \text{if } I_j < \theta \\ w_j & \text{otherwise} \end{cases}$
这里 $w_j'$ 是剪枝后的权重。

4.重训练

再训练模型：在剪枝后进行再训练，以调整剩余的权重，使模型恢复性能。可以使用梯度下降优化：
$w_j \leftarrow w_j - \eta g_j$
其中 $η\eta$ 是学习率。

基于CNN的梯度重要性剪枝代码

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms

# 定义简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=5)  # 输入通道改为3（RGB）
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        self.fc1 = nn.Linear(32 * 5 * 5, 128)  # 根据输入尺寸调整
        self.fc2 = nn.Linear(128, 10)  # 输出10个类

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 32 * 5 * 5)  # 根据输入尺寸调整
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 训练模型
def train_model(model, train_loader, optimizer, criterion, num_epochs=1):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)  # 将数据转移到GPU
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()  # 累加损失

        average_loss = total_loss / len(train_loader)  # 计算平均损失
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {average_loss:.4f}')  # 输出平均损失

# 剪枝权重
def prune_weights(model, threshold):
    with torch.no_grad():
        for name, param in model.named_parameters():
            if 'weight' in name:  # 只剪枝权重参数
                importance = param.grad.abs()  # 计算权重的梯度绝对值
                mask = importance < threshold  # 创建剪枝掩码
                param[mask] = 0  # 剪除不重要的权重

# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 选择设备：如果有可用的MPS，使用MPS；否则如果有可用的GPU，使用GPU；否则使用CPU
if torch.backends.mps.is_available():
    device = torch.device("mps")  # MacBook M3
elif torch.cuda.is_available():
    device = torch.device("cuda")  # Windows GPU
else:
    device = torch.device("cpu")  # CPU

# 初始化模型、损失函数和优化器
model = SimpleCNN().to(device)  # 将模型转移到GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
train_model(model, train_loader, optimizer, criterion, num_epochs=5)

# 计算梯度并剪枝
model.train()  # 确保模型处于训练模式
for data, target in train_loader:
    data, target = data.to(device), target.to(device)  # 将数据转移到GPU
    optimizer.zero_grad()  # 清空之前的梯度
    output = model(data)
    loss = criterion(output, target)
    loss.backward()  # 计算梯度
    break  # 只计算一次梯度用于剪枝

# 剪枝阈值设定
threshold = 0.01  # 设定剪枝阈值
prune_weights(model, threshold)

# 重训练模型
train_model(model, train_loader, optimizer, criterion, num_epochs=5)

print("模型剪枝与重训练完成。")

注意

置为零的权重的影响

前向传播：
- 在前向传播阶段，设置为零的权重确实不会影响模型的输出。具体来说，如果权重 $w_j=0$ ，那么在计算输出时，所有与 $w_j$ 相关的激活值都会被乘以零，结果为零。因此，零权重不会对损失函数产生任何贡献。
损失函数的贡献：
- 损失函数是根据模型输出和真实标签之间的差异计算的。对于那些与零权重相关的部分，它们的输出为零，所以这些权重不会影响损失的计算，也就不会产生梯度（贡献）。
反向传播：
- 在反向传播过程中，只有对损失函数有贡献的权重才会计算梯度。由于零权重没有参与损失函数的计算，它们的梯度将会是零。

基于LLAMA2梯度重要性剪枝的代码

import torch
from transformers import LlamaForCausalLM, LlamaTokenizer
from transformers import AdamW
from datasets import load_dataset

def gradient_based_pruning(model, pruning_rate):
    """
    对模型进行梯度剪枝
    :param model: 待剪枝的模型
    :param pruning_rate: 剪枝率 (0到1之间的浮点数)
    """
    for name, param in model.named_parameters():
        if 'weight' in name and param.grad is not None:
            # 计算权重的梯度绝对值
            weight_grad = param.grad.abs()
            
            # 计算剪枝阈值
            threshold = torch.quantile(weight_grad, pruning_rate)
            
            # 剪枝
            mask = weight_grad < threshold
            param.data[mask] = 0  # 将低于阈值的权重置为0

# 加载IMDB情感分析数据集
dataset = load_dataset("imdb")

# 加载预训练的LLAMA2模型和分词器
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b")
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b")

# 准备输入数据（取前1000条样本进行示例）
train_texts = dataset['train']['text'][:1000]
train_labels = dataset['train']['label'][:1000]

# 编码输入数据
inputs = tokenizer(train_texts, padding=True, truncation=True, return_tensors="pt")

# 设置优化器
optimizer = AdamW(model.parameters(), lr=5e-5)

# 训练循环示例
model.train()
for epoch in range(3):  # 假设训练3个epoch
    optimizer.zero_grad()  # 清除梯度
    
    # 获取模型输出
    outputs = model(input_ids=inputs['input_ids'], labels=inputs['input_ids'])
    loss = outputs.loss  # 获取损失
    loss.backward()  # 反向传播

    # 在这里进行梯度剪枝
    pruning_rate = 0.2  # 剪去20%的权重
    gradient_based_pruning(model, pruning_rate)

    optimizer.step()  # 更新模型参数

    print(f"Epoch {epoch + 1}, Loss: {loss.item()}")

# 训练完成后，你可以继续评估剪枝后的模型