1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
| from transformers import ( BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments, ) import torch from datasets import load_dataset, DatasetDict
"""## 1.加载和准备IMDB数据集样本 选取一部分数据用于Fine-tuning。 """
dataset = load_dataset('imdb', split='train') small_dataset = dataset.shuffle(seed=42).select(range(10000))
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
device = "cuda" if torch.cuda.is_available() else "cpu"
def encode(examples): return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
encoded_small_dataset = small_dataset.map(encode, batched=True)
import torch from transformers import BertModel import matplotlib.pyplot as plt
def visualize_attention(sentence, model, tokenizer): model.to(device) model.eval()
inputs = tokenizer(sentence, return_tensors="pt").to(device)
with torch.no_grad(): outputs = model(**inputs) attentions = outputs.attentions
layer = 5 head = 1 attention = attentions[layer][0, head].cpu().numpy()
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0].cpu()) plt.figure(figsize=(10, 10)) plt.matshow(attention, cmap='viridis') plt.xticks(range(len(tokens)), tokens, rotation=90) plt.yticks(range(len(tokens)), tokens) plt.colorbar() plt.show()
"""## 2. 可视化一个样本句子的注意力权重(未经Fine-tuning) 选择数据集中的一个句子并展示其原始BERT模型的注意力权重。 """
model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True) sample_sentence = "I love this movie, it's fantastic!" visualize_attention(sample_sentence, model, tokenizer)
"""## 3. Fine-tuning BERT模型 在选取的IMDB样本上进行Fine-tuning。 ### 3.1 准备数据加载器 为了训练模型,我们需要创建PyTorch的DataLoader。这将使我们能够在训练过程中有效地加载数据。
### 3.2 设置Fine-tuning环境 初始化模型、优化器以及损失函数。
### 3.3 Fine-tuning模型 执行Fine-tuning的训练循环。执行以上代码将在IMDB数据集的小样本上对BERT模型进行Fine-tuning。这可能需要一些时间,具体取决于您的硬件配置。
"""
from torch.utils.data import DataLoader
encoded_small_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])
train_loader = DataLoader(encoded_small_dataset, batch_size=8, shuffle=True)
from transformers import BertConfig, BertForSequenceClassification import torch.optim as optim
config = BertConfig.from_pretrained("bert-base-uncased", output_attentions=True)
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", config=config)
optimizer = optim.AdamW(model.parameters(), lr=5e-5)
criterion = torch.nn.CrossEntropyLoss()
model.to(device)
epochs = 8
for epoch in range(epochs): model.train() total_loss = 0 for batch in train_loader: input_ids = batch['input_ids'].to(device) attention_mask = batch['attention_mask'].to(device) labels = batch['label'].to(device)
outputs = model(input_ids, attention_mask=attention_mask)
loss = criterion(outputs.logits, labels)
optimizer.zero_grad() loss.backward() optimizer.step()
total_loss += loss.item()
print(f"Epoch: {epoch+1}, Loss: {total_loss/len(train_loader)}")
"""### 4. 可视化同一句子的注意力权重(经过Fine-tuning) 使用Fine-tuning后的模型再次可视化同一句子的注意力权重。您可以重用之前提供的visualize_attention函数: """
visualize_attention(sample_sentence, model, tokenizer)
|