基于 Transformer RoBERTa的情感分类任务实践总结之五——剪枝

注意力头重要性剪枝（Head Importance Pruning）

这种剪枝方法旨在从像 RoBERTa 这样的预训练 Transformer 模型中移除最不重要的注意力头，同时不显著降低其性能。其核心思想是量化每个注意力头对模型整体损失的贡献，然后移除贡献最小的那些头。

算法总结

1. 初始化注意力头掩码（Head Mask）：

   • 创建一个 head_mask 张量，并设置 requires_grad=True。这个掩码的维度与模型中层数和每层注意力头的数量相同。
   • 初始时，掩码中的所有值都设为 1，这意味着所有注意力头都是活跃的。

2. 计算注意力头重要性：

   • 将模型设置为评估模式 (model.eval())。
   • 对于给定的一批数据，执行模型的正向传播，并应用 head_mask。head_mask 作为注意力头输出的乘数，让模型“看到”哪些头正在被考虑，是权重调节器。
   • 计算损失。
   • 关键步骤是调用 loss.backward()。这一步会计算损失对 head_mask 的梯度。这些梯度的绝对值可以作为每个头重要性的代理——更大的梯度表明改变该头的输出（通过掩码）对损失有更大的影响，从而暗示其重要性更高。
   • 梯度被分离（detach）并按层内注意力头进行归一化，以获得相对重要性分数。

3. 识别要剪枝的注意力头：

   • 对于每一层，识别出重要性分数最低（由计算出的梯度决定）的注意力头。
   • 将这些最不重要的头收集到一个字典 (heads_to_prune) 中，该字典将层索引映射到要剪枝的头索引列表。

4. 剪枝注意力头：

   • 调用 model.prune_heads() 方法，并传入 heads_to_prune 字典。该方法会内部修改模型的架构，以移除指定的注意力头。这通常涉及调整后续层（例如输出投影层）的维度，以适应减少的注意力头数量。

5. 微调：

   • 剪枝后，通常建议在您的数据集上对剪枝后的模型进行微调。这使得剩余的头和模型其他参数能够适应剪枝后的架构，并可能恢复任何性能损失。在您的代码中，模型在剪枝后立即被传递给 Trainer 进行训练。

为什么它有效？

这种方法之所以有效，是因为它直接量化了每个注意力头对模型主要目标（最小化损失）的贡献。通过移除贡献最小的头，实际上是在降低模型的冗余和计算复杂性，同时旨在保留其最关键的信息通路。

代码

import torch
from transformers import (
    RobertaForSequenceClassification,
    RobertaTokenizer,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset
import os
from huggingface_hub import snapshot_download
from torch.utils.data import DataLoader
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 下载模型到本地
model_path = snapshot_download(
    repo_id="roberta-base",
    local_dir="./local_roberta_model",
    local_dir_use_symlinks=False,
)
print(f"模型已下载到: {model_path}")

# 加载 tokenizer & 模型
tokenizer = RobertaTokenizer.from_pretrained(model_path,   
                                             local_files_only=True  # 关键参数：仅使用本地文件，找不到时抛出错误
)
model = RobertaForSequenceClassification.from_pretrained(model_path, num_labels=2)

# 加载 IMDB & 预处理
dataset = load_dataset("imdb")
small_train = dataset["train"].shuffle(seed=42)#.select(range(1000))
small_test = dataset["test"].shuffle(seed=42)#.select(range(1000))

def preprocess_function(examples):
    return tokenizer(
        examples["text"], truncation=True, padding="max_length", max_length=128
    )

small_train = small_train.map(preprocess_function, batched=True)
small_test = small_test.map(preprocess_function, batched=True)
#  Hugging Face Trainer 默认期望这个目标列的名称是 "labels"。
small_train = small_train.rename_column("label", "labels")
small_test = small_test.rename_column("label", "labels")

small_train.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
small_test.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

train_loader = DataLoader(small_train, batch_size=8)

# Head Importance via head_mask
def compute_head_importance(model, dataloader):
    model.eval()
    head_mask = torch.ones(
        model.config.num_hidden_layers,
        model.config.num_attention_heads,
        requires_grad=True,
        device=model.device,
    )

    for batch in dataloader:
        batch = {k: v.to(model.device) for k, v in batch.items()}
        outputs = model(**batch, head_mask=head_mask)
        loss = outputs.loss
        loss.backward()

    head_importance = head_mask.grad.abs().detach()
    head_importance /= head_importance.sum(dim=-1, keepdim=True)
    return head_importance

# 计算并剪枝
head_importance = compute_head_importance(model, train_loader)
print("Head Importance:\n", head_importance)

heads_to_prune = {}
for layer in range(head_importance.size(0)):
    least_important = torch.argmin(head_importance[layer]).item()
    heads_to_prune[layer] = [least_important]

print("Heads to prune:\n", heads_to_prune)
model.prune_heads(heads_to_prune)

# 定义评价指标
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy_score(labels, predictions)
    f1 = f1_score(labels, predictions)
    return {"accuracy": acc, "f1": f1}

# Trainer 参数 & 日志
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=2,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    learning_rate=2e-5,
    eval_strategy="epoch",  # 每个 epoch 评估一次
    logging_dir="./logs",#tensorboard --logdir ./logs
    logging_steps=20,
    report_to=["tensorboard"],
    save_strategy="epoch",  # 每个 epoch 保存
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train,
    eval_dataset=small_test,
    processing_class=tokenizer,
    compute_metrics=compute_metrics,
)

# 训练 & 评估
trainer.train()
metrics = trainer.evaluate()
print("Final Evaluation Metrics:", metrics)

评价

TensorBoard