基于 Transformer RoBERTa的情感分类任务实践总结之五——剪枝
注意力头重要性剪枝（Head Importance Pruning）
这种剪枝方法旨在从像 RoBERTa 这样的预训练 Transformer 模型中移除最不重要的注意力头，同时不显著降低其性能。其核心思想是量化每个注意力头对模型整体损失的贡献，然后移除贡献最小的那些头。
算法总结
-
初始化注意力头掩码（Head Mask）：
- 创建一个
head_mask张量，并设置requires_grad=True。这个掩码的维度与模型中层数和每层注意力头的数量相同。 - 初始时，掩码中的所有值都设为 1，这意味着所有注意力头都是活跃的。
- 创建一个
-
计算注意力头重要性：
- 将模型设置为评估模式 (
model.eval())。 - 对于给定的一批数据，执行模型的正向传播，并应用
head_mask。head_mask作为注意力头输出的乘数，让模型“看到”哪些头正在被考虑，是权重调节器。 - 计算损失。
- 关键步骤是调用
loss.backward()。这一步会计算损失对head_mask的梯度。这些梯度的绝对值可以作为每个头重要性的代理——更大的梯度表明改变该头的输出（通过掩码）对损失有更大的影响，从而暗示其重要性更高。 - 梯度被分离（detach）并按层内注意力头进行归一化，以获得相对重要性分数。
- 将模型设置为评估模式 (
-
识别要剪枝的注意力头：
- 对于每一层，识别出重要性分数最低（由计算出的梯度决定）的注意力头。
- 将这些最不重要的头收集到一个字典 (
heads_to_prune) 中，该字典将层索引映射到要剪枝的头索引列表。
-
剪枝注意力头：
- 调用
model.prune_heads()方法，并传入heads_to_prune字典。该方法会内部修改模型的架构，以移除指定的注意力头。这通常涉及调整后续层（例如输出投影层）的维度，以适应减少的注意力头数量。
- 调用
-
微调：
- 剪枝后，通常建议在您的数据集上对剪枝后的模型进行微调。这使得剩余的头和模型其他参数能够适应剪枝后的架构，并可能恢复任何性能损失。在您的代码中，模型在剪枝后立即被传递给
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_scoreos.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.lossloss.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_predpredictions = 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 ./logslogging_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
