我们用Bert实现对THUCNews文本分类,在embdedding使用FGSM添加干扰实现对抗训练,并且考虑对抗性样本的防御,代码实现并进行逐行注释,使用pytorch实现。
使用BERT实现文本分类的步骤
- 安装相应的依赖项(如pytorch、transformers等)。
- 准备THUCNews数据集,并将其划分为训练集、验证集和测试集。
- 加载预训练的BERT模型。
- 训练模型,使用交叉熵损失函数和Adam优化器。
- 预测结果并评估模型的性能。
在这个过程中,我们还需要考虑如何使用FGSM添加干扰实现对抗训练,并进行对抗性样本的防御。接下来,我们将逐步地介绍如何实现这些功能。
目录
一、准备数据
首先是数据准备阶段。THUCNews是一个中文文本分类数据集,包含14个类别的新闻文章。可以从这里下载数据集。我们可以使用Python的pandas库来读取数据并将其划分为训练集、验证集和测试集。
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
df = pd.read_csv("path/to/thucnews.csv")
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42) # 划分训练集和测试机
train_df, validate_df = train_test_split(train_df, test_size=0.2, random_state=42) # 划分训练集和验证集
train_texts = train_df["text"].tolist()
train_labels = train_df["label"].tolist()
validate_texts = validate_df["text"].tolist()
validate_labels = validate_df["label"].tolist()
test_texts = test_df["text"].tolist()
test_labels = test_df["label"].tolist()
二、加载BERT模型
接下来,我们需要加载预训练的BERT模型。PyTorch中已经提供了一些预训练的BERT模型,我们可以使用transformers库来加载它们。
from transformers import BertTokenizer, BertForSequenceClassification
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese") # 加载中文BERT模型的tokenizer
model = BertForSequenceClassification.from_pretrained("bert-base-chinese", num_labels=14) # 加载中文BERT模型
在这段代码中,我们使用from_pretrained
函数来加载BERT模型,并指定分类数目为14(THUCNews数据集的类别数目)。注意,我们还需要使用相应的tokenizer来对文本进行tokenization,因为BERT模型需要文本表示为token的形式。
三、训练模型
我们可以使用PyTorch内置的DataLoader
来加载数据,并使用交叉熵损失函数和Adam优化器来训练模型。
import torch
from torch.utils.data import DataLoader, Dataset
class TextDataset(Dataset):
def __init__(self, texts, labels, tokenizer):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = self.labels[idx]
inputs = self.tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=256,
padding='max_length',
return_attention_mask=True,
return_tensors='pt'
)
return {
'input_ids': inputs['input_ids'][0],
'attention_mask': inputs['attention_mask'][0],
'labels': torch.tensor(label, dtype=torch.long)
}
train_dataset = TextDataset(train_texts, train_labels, tokenizer)
validate_dataset = TextDataset(validate_texts, validate_labels, tokenizer)
test_dataset = TextDataset(test_texts, test_labels, tokenizer)
train_dataloader = DataLoader(train_dataset, batch_size=32)
validate_dataloader = DataLoader(validate_dataset, batch_size=32)
test_dataloader = DataLoader(test_dataset, batch_size=32)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
for epoch in range(10):
running_loss = 0.0
for data in train_dataloader:
input_ids = data['input_ids'].to(device)
attention_mask = data['attention_mask'].to(device)
labels = data['labels'].to(device)
optimizer.zero_grad()
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs[0]
loss.backward()
optimizer.step()
running_loss += loss.item()
print('Epoch:', epoch, 'Loss:', running_loss / len(train_dataloader))
在每个epoch中,我们使用train_dataloader分批加载训练数据,并使用optimizer对模型进行优化。交叉熵损失函数的定义可以参考这里。
四、添加对抗性干扰
接下来,我们将用FGSM方法添加对抗性干扰。我们可以定义一个函数,该函数将接受文本和标签,使用BERT模型进行前向传播,并在embedding层上添加一定的干扰,最后返回对抗性样本及其标签。具体实现如下:
def fgsm_attack(model, criterion, input_ids, attention_mask, labels, epsilon):
input_ids.requires_grad = True
attention_mask.requires_grad = True
outputs = model(input_ids, attention_mask=attention_mask)
loss = criterion(outputs.logits, labels)
model.zero_grad()
loss.backward()
input_ids_grad = input_ids.grad.detach()
attention_mask_grad = attention_mask.grad.detach()
sign_data_grad = input_ids_grad.sign()
perturbed_input_ids = input_ids + epsilon * sign_data_grad
sign_attention_grad = attention_mask_grad.sign()
perturbed_attention_mask = attention_mask + epsilon * sign_attention_grad
return perturbed_input_ids, perturbed_attention_mask, labels
在上面的代码中,我们首先将文本和attention mask设置为可求导的,并使用该模型对其进行前向传播和计算损失。然后,我们计算输入的梯度值,并使用sign函数得到攻击向量。最后,我们将攻击向量添加到原始输入上,并返回扰动后的输入。
五、防御对抗性样本
如果我们使用FGSM方法生成对抗性样本,那么输入会发生变化,从而导致模型性能下降。为了解决这个问题,我们可以使用预训练的中文RoBERTa模型RoBERTa-wwm-ext和Adversarial Training来防御对抗性样本。
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification
roberta_model = AutoModelForSequenceClassification.from_pretrained("model_path", num_labels=14).cuda()
def adv_training(model, criterion, train_dataloader, validate_dataloader, optimizer, num_epochs, epsilon):
model.train()
for epoch in range(num_epochs):
running_loss = 0.0
for data in train_dataloader:
input_ids = data['input_ids'].to(device)
attention_mask = data['attention_mask'].to(device)
labels = data['labels'].to(device)
perturbed_input_ids, perturbed_attention_mask, labels = \
fgsm_attack(model, criterion, input_ids, attention_mask, labels, epsilon)
model.zero_grad()
outputs = model(perturbed_input_ids, attention_mask=perturbed_attention_mask, labels=labels)
loss = outputs[0]
loss.backward()
optimizer.step()
running_loss += loss.item()
print('Epoch:', epoch, 'Loss:', running_loss / len(train_dataloader))
correct = 0
total = 0
model.eval()
with torch.no_grad():
for data in validate_dataloader:
input_ids = data['input_ids'].to(device)
attention_mask = data['attention_mask'].to(device)
labels = data['labels'].to(device)
outputs = model(input_ids, attention_mask=attention_mask)
_, predicted = torch.max(F.softmax(outputs.logits, dim=1), 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy on validation set:', 100 * correct / total)
model.train()
return model
adv_model = adv_training(model, criterion, train_dataloader, validate_dataloader, optimizer, 10, 0.5)
在这个防御方案中,我们首先使用fgsm_attack
函数生成对抗性样本,然后使用预训练的中文RoBERTa模型进行防御。具体地,我们将其作为一个特殊的提取器来提取原始文本的embedding,并在这个embedding层上添加对抗性扰动。最后,我们训练一个文本分类模型,使用这些扰动后的embedding来提高模型的鲁棒性。
到这里,我们就完成了使用BERT模型实现文本分类,并添加了对抗性扰动来实现对抗训练和防御对抗性样本的实验。