Bert+对抗训练+对抗性样本防御的文本分类实战

我们用Bert实现对THUCNews文本分类,在embdedding使用FGSM添加干扰实现对抗训练,并且考虑对抗性样本的防御,代码实现并进行逐行注释,使用pytorch实现。

使用BERT实现文本分类的步骤

  1. 安装相应的依赖项(如pytorch、transformers等)。
  2. 准备THUCNews数据集,并将其划分为训练集、验证集和测试集。
  3. 加载预训练的BERT模型。
  4. 训练模型,使用交叉熵损失函数和Adam优化器。
  5. 预测结果并评估模型的性能。

在这个过程中,我们还需要考虑如何使用FGSM添加干扰实现对抗训练,并进行对抗性样本的防御。接下来,我们将逐步地介绍如何实现这些功能。

目录

使用BERT实现文本分类的步骤

一、准备数据

二、加载BERT模型

三、训练模型 

四、添加对抗性干扰 

五、防御对抗性样本 


一、准备数据

首先是数据准备阶段。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模型实现文本分类,并添加了对抗性扰动来实现对抗训练和防御对抗性样本的实验。

猜你喜欢

转载自blog.csdn.net/weixin_43734080/article/details/130900946