PGD实现中文文本分类（Loss=0.5L1+0.5L2)网站首页 学无止境

Bert+FGSM/PGD实现中文文本分类（Loss=0.5L1+0.5L2)

Dr.sky_ 2024-08-17 00:01:02

简介Bert+FGSM/PGD实现中文文本分类（Loss=0.5L1+0.5L2)

任务目标：在使用FGSM/PGD来训练Bert模型进行文本分类，其实现原理可以简单概括为以下几个步骤：

对原始文本每个词转换为对应的嵌入向量。
将每个嵌入向量与一个小的扰动向量相加，从而生成对抗样本。这个扰动向量的大小可以通过一个超参数来控制。
将生成的对抗样本和原始样本一起用于训练模型。具体来说，可以将它们组成一个batch，然后使用交叉熵损失函数来训练模型。
在训练过程中，可以周期性地增加扰动向量的大小，从而使得模型逐渐适应更强的攻击。这个过程可以称为“逐步增强对抗性训练”。
通过使用FGSM/PGD来训练Bert模型，可以使得模型对对抗样本更加鲁棒，从而提高其在真实场景中的泛化能力和分类准确率。
在训练过程中我们设置 总样本Loss=0.5原样本Loss+0.5对抗样本Loss，来提升模型的鲁棒性。

在使用FGSM/PGD来训练Bert模型进行文本分类时，其实现原理可以概括为以下几个步骤：

一、导入所需的库和模块

这段代码主要是导入了一些必要的 PyTorch 和 transformers 库中的类和函数，其中： torch 是 PyTorch 库的主要模块，包含了大量的张量操作和神经网络模块等。 nn 是 PyTorch 中的神经网络模块，包含了各种神经网络层和模型等。 optim 是 PyTorch 中的优化器模块，包含了各种优化算法，如 SGD、Adam 等。 DataLoader 和 Dataset 是 PyTorch 中的数据集和数据加载器模块，用于加载和处理数据集。 BertTokenizerFast 是 transformers 库中的类，用于将文本转换为 BERT 模型的输入格式。 BertForSequenceClassification 是 transformers 库中的类，用于进行文本分类任务。

# 导入 PyTorch 库
import torch
# 导入 PyTorch 中的神经网络模块
import torch.nn as nn
# 导入 PyTorch 中的优化器模块
import torch.optim as optim
# 导入 PyTorch 中的数据集和数据加载器模块
from torch.utils.data import DataLoader, Dataset
# 导入 transformers 库中的 BertTokenizerFast 和 BertForSequenceClassification 类
from transformers import BertTokenizerFast, BertForSequenceClassification
import numpy as np

二、加载数据集

这段代码的主要作用是创建一个用于加载 THUCNews 数据集的数据集类 THUCNewsDataset，并实现 len 和 getitem 方法。其中： tqdm 库用于显示进度条，可以让我们在读取数据集时更直观地了解进度。 BertTokenizerFast.from_pretrained('bert-base-chinese') 创建了一个 BertTokenizerFast 对象，用于将文本转换为 BERT 模型的输入格式。 self.data 列表用于存储数据集中的每个样本，每个样本是一个元组，包含文本和标签。 init 方法用于初始化数据集对象。在该方法中，我们打开数据集文件，并逐行读取数据。对于每一行数据，我们使用 strip() 方法去除空格和换行符，然后使用 split(' ') 方法将文本和标签分开。最后，我们将文本和标签封装成一个元组，并将其添加到 self.data 列表中。 len 方法用于返回数据集的长度，即数据集中样本的个数。 getitem 方法用于获取数据集中的一个样本。在该方法中，我们从 self.data 列表中获取第 idx 个样本的文本和标签。然后，我们使用 self.tokenizer 将文本转换为 BERT 模型的输入格式，并将标签转换为 tensor，并将其添加到 inputs 字典中。最后，我们返回 inputs 字典。

# 导入 tqdm 库
from tqdm import tqdm
# 定义一个 THUCNewsDataset 类，继承自 PyTorch 中的 Dataset 类
class THUCNewsDataset(Dataset):
    # 定义构造函数，接收一个文件路径作为参数
    def __init__(self, file_path):
        # 初始化 BERT tokenizer
        self.tokenizer = BertTokenizerFast.from_pretrained('bert-base-chinese')
        # 定义一个列表，用于存储数据集中的每一条数据
        self.data = []
        # 打开数据集文件，逐行读取数据并处理
        with open(file_path, 'r', encoding='utf-8') as f:
            # 使用 tqdm 库显示读取进度
            for line in tqdm(f):
                # 从每一行数据中提取文本和标签，并将其存储到列表中
                text, label = line.strip().split('	')
                self.data.append((text, int(label)))
    
    # 定义 __len__ 方法，返回数据集的大小
    def __len__(self):
        return len(self.data)
    
    # 定义 __getitem__ 方法，根据索引返回数据集中的一条数据
    def __getitem__(self, idx):
        # 从列表中获取文本和标签
        text, label = self.data[idx]
        # 使用 BERT tokenizer 对文本进行处理，将其转换为 BERT 模型的输入格式
        inputs = self.tokenizer(text, padding='max_length', truncation=True, max_length=32, return_tensors='pt')
        # 将标签转换为 PyTorch 的张量格式，并将其添加到输入中
        inputs['labels'] = torch.tensor(label)
        # 返回处理后的输入
        return inputs

# 加载训练集、测试集和验证集
train_dataset = THUCNewsDataset('train.txt')
test_dataset = THUCNewsDataset('test.txt')
dev_dataset = THUCNewsDataset('dev.txt')


# 导入 PyTorch 库中的 pad_sequence 函数，用于填充序列
from torch.nn.utils.rnn import pad_sequence

# 定义一个 collate_fn 函数，用于对数据进行批处理
def collate_fn(batch):
    # 从批次数据中提取 input_ids、attention_mask 和 labels
    input_ids = [item['input_ids'] for item in batch]
    attention_mask = [item['attention_mask'] for item in batch]
    labels = [item['labels'] for item in batch]
    # 对 input_ids 和 attention_mask 进行填充操作，使它们的长度相同
    input_ids = pad_sequence(input_ids, batch_first=True, padding_value=0)
    attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0)
    # 将 labels 转换为 tensor 类型
    labels = torch.tensor(labels)
    # 返回一个字典，包含处理后的 input_ids、attention_mask 和 labels
    return {'input_ids': input_ids, 'attention_mask': attention_mask, 'labels': labels}

# 创建数据加载器，用于批量加载数据
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)
dev_loader = DataLoader(dev_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)

三、定义模型和优化器

这段代码的主要作用是创建一个用于文本分类的 BERT 模型，并初始化优化器和损失函数。其中： BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=10) 加载了预训练的 BERT 模型，并创建了一个用于文本分类的 BERT 模型。其中，'bert-base-chinese' 表示使用中文 BERT 模型，num_labels=10 表示模型的输出类别数为 10。 optim.Adam(model.parameters(), lr=2e-5) 创建了一个 Adam 优化器，用于更新模型参数。其中，model.parameters() 表示优化器需要更新的模型参数，lr=2e-5 表示学习率为 2e-5。 nn.CrossEntropyLoss() 创建了一个交叉熵损失函数，用于计算模型的损失。在文本分类任务中，通常使用交叉熵损失函数作为损失函数。

# 加载预训练的 BERT 模型，并创建一个用于文本分类的 BERT 模型
model = BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=10)
# 创建一个 Adam 优化器，用于更新模型参数
optimizer = optim.Adam(model.parameters(), lr=2e-5)
# 创建一个交叉熵损失函数，用于计算模型的损失
criterion = nn.CrossEntropyLoss()

上面红色的提示是正常的，大概意思是指Bert用在下游任务需要微调。

四、基于原生Bert文本分类

4.1 定义训练函数

训练函数用于训练模型，测试函数用于测试模型在测试数据集上的性能。

# 训练函数在每个批次中都进行了反向传播和参数更新
def train(model, optimizer, criterion, train_loader, device):
    model.train() # 将模型设置为训练模式
    train_loss = 0 # 初始化训练损失为0
    train_acc = 0 # 初始化训练准确率为0
    for batch in train_loader: # 遍历训练数据集
        input_ids = batch['input_ids'].squeeze(1).to(device) # 将输入数据移动到计算设备上
        attention_mask = batch['attention_mask'].squeeze(1).to(device) # 将输入数据移动到计算设备上
        labels = batch['labels'].to(device) # 将标签移动到计算设备上
        optimizer.zero_grad() # 清空梯度
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels) # 模型前向传播
        loss = criterion(outputs.logits, labels) # 计算损失
        train_loss += loss.item() # 累加损失
        loss.backward() # 反向传播，计算梯度
        optimizer.step() # 更新参数
        preds = torch.argmax(outputs.logits, dim=1) # 计算预测结果
        train_acc += torch.sum(preds == labels).item() # 计算准确率
    train_loss /= len(train_loader) # 计算平均损失
    train_acc /= len(train_loader.dataset) # 计算平均准确率
    return train_loss, train_acc # 返回训练损失和准确率
# 验证函数只进行了前向传播，没有进行反向传播和参数更新
def evaluate(model, criterion, test_loader, device):
    model.eval() # 将模型设置为评估模式
    test_loss = 0 # 初始化测试损失为0
    test_acc = 0 # 初始化测试准确率为0
    with torch.no_grad(): # 关闭梯度计算
        for batch in test_loader: # 遍历测试数据集
            input_ids = batch['input_ids'].squeeze(1).to(device) # 将输入数据移动到计算设备上
            attention_mask = batch['attention_mask'].squeeze(1).to(device) # 将输入数据移动到计算设备上
            labels = batch['labels'].to(device) # 将标签移动到计算设备上
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels) # 模型前向传播
            loss = criterion(outputs.logits, labels) # 计算损失
            test_loss += loss.item() # 累加损失
            preds = torch.argmax(outputs.logits, dim=1) # 计算预测结果
            test_acc += torch.sum(preds == labels).item() # 计算准确率
    test_loss /= len(test_loader) # 计算平均损失
    test_acc /= len(test_loader.dataset) # 计算平均准确率
    return test_loss, test_acc # 返回测试损失和准确率

4.2 定义测试函数

# 将模型移动到计算设备上（GPU 或 CPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 训练模型
best_acc = 0 # 初始化最佳准确率为0
for epoch in range(10): # 遍历10个 epoch
    train_loss, train_acc = train(model, optimizer, criterion, train_loader, device) # 训练模型
    test_loss, test_acc = evaluate(model, criterion, test_loader, device) # 在测试集上评估模型
    dev_loss, dev_acc = evaluate(model, criterion, dev_loader, device) # 在验证集上评估模型
    # 输出当前 epoch 的训练损失、训练准确率、测试损失、测试准确率、验证损失和验证准确率
    print(f'Epoch {epoch + 1}: Train Loss {train_loss:.4f}, Train Acc {train_acc:.4f}, Test Loss {test_loss:.4f}, Test Acc {test_acc:.4f}, Dev Loss {dev_loss:.4f}, Dev Acc {dev_acc:.4f}')
    if dev_acc > best_acc: # 如果当前 epoch 的验证准确率大于历史最佳准确率
        best_acc = dev_acc # 更新历史最佳准确率
        torch.save(model.state_dict(), 'best_model.pt') # 保存模型参数到文件 best_model.pt

五、Bert+FGSM文本分类

以下是使用FGSM在embedding添加干扰，并考虑到对抗性样本的防御的训练和测试函数。我们在train()函数中添加了epsilon参数，以控制对抗性样本的干扰程度。我们还计算了原始样本的损失和对抗性样本的损失，并将它们加权平均作为总损失。在evaluate()函数中，我们仅对模型进行前向传递，以便计算测试损失和准确率。 BertTokenizerFast 是 transformers 库中的一个高速分词器，它是 BertTokenizer 的改进版本。与 BertTokenizer 不同，BertTokenizerFast 使用 Rust 实现的底层代码，因此在分词速度方面更快。另外，BertTokenizerFast 还支持更多的特殊标记，例如 Truncation，Padding，以及更好的处理未知单词（out-of-vocabulary，OOV）。如果你的代码需要处理大量文本数据，那么使用 BertTokenizerFast 可以显著提高代码的执行效率。但是，如果你的代码只需要处理少量文本数据，那么使用 BertTokenizer 更加方便和易于使用。

5.1 定义FGSM对抗训练函数

这是一个用于生成对抗样本的函数，输入参数包括原始文本嵌入向量embedding，扰动大小epsilon和梯度gradient。该函数会计算梯度的符号，创建扰动，并将扰动限制在有效范围内。最后返回生成的对抗样本。这段代码实现了 FGSM 对抗攻击，目的是在原始输入的嵌入（embeddings）中添加一些干扰，以生成对抗样本（adversarial sample）。

具体来说，这个函数接受三个参数：原始输入的嵌入（embeddings）、对于原始输入的梯度（grad）、添加干扰的程度（epsilon）。它首先将梯度值（grad）符号化（sign），得到输入的梯度符号，然后与干扰程度（epsilon）相乘，得到干扰值（perturb）。最后，将干扰（perturb）添加到原始输入的嵌入中，得到对抗样本的嵌入（perturb_embeds）。

FGSM 对抗攻击的核心思想是在保证对抗样本与原始样本之间尽可能小的距离（即干扰程度尽可能小）的同时，使得对抗样本能够欺骗深度学习模型。这个距离被称为 L_p 距离，通常选择 L_∞ 距离，即干扰程度的上界为 epsilon。

def fgsm_attack(embedding, epsilon, gradient):
    # 计算梯度的符号
    if gradient is None:
        print('gradient is None')
        return embedding
    sign_gradient = gradient.sign()
    # 创建扰动
    perturbed_embedding = embedding + epsilon * sign_gradient
    # 将扰动限制在有效范围内
    perturbed_embedding = torch.clamp(perturbed_embedding, min=0, max=1)
    return perturbed_embedding

5.2 定义训练模型函数

这是用于训练模型的函数，输入参数包括模型model，优化器optimizer，损失函数criterion，训练数据集的数据加载器train_loader，计算设备device和。该函数会遍历训练数据集，将输入数据和标签移动到计算设备上，清空梯度，生成对抗样本，计算损失和梯度，更新参数，计算准确率等操作，并返回训练损失和准确率。

学习率调度器的作用是在训练过程中自动调整学习率，以提高模型的训练效果。 ReduceLROnPlateau 是一个 PyTorch 自带的学习率调度器类，它有以下参数： optimizer：优化器对象，用于更新模型参数； mode：模式，可选值为 min、max 或 auto，表示监测的指标是越小越好、越大越好还是自动选择。这里我们选择 max，表示准确率越大越好； factor：学习率缩放因子，每次调整学习率时将当前学习率乘以该因子； patience：当监测指标在 patience 轮内没有变化时，减小学习率； verbose：是否打印调度器信息； epsilon：学习率变化的最小阈值，如果新学习率与旧学习率之间的差异小于该阈值，则不会更新学习率； cooldown：调整学习率后，暂停更新学习率的轮数； min_lr：学习率的下限； eps：数值稳定性参数。当调用 scheduler.step(acc) 时，调度器会根据当前的准确率 acc 来自动调整学习率。如果在 patience 轮内准确率没有提高，则会将学习率缩小 factor 倍，直到学习率达到下限 min_lr。

# 创建一个 Adam 优化器，用于更新模型参数
parameters = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.Adam(parameters, lr=2e-5)
scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3, verbose=1, epsilon=1e-4, cooldown=0, min_lr=0, eps=1e-8)
def train(model, optimizer, criterion, train_loader, device, epsilon):
    model.train() # 将模型设置为训练模式
    train_loss = 0 # 初始化训练损失为0
    train_acc = 0 # 初始化训练准确率为0
    tokenizer = BertTokenizerFast.from_pretrained('bert-base-chinese')
    for batch in train_loader: # 遍历训练数据集
        input_ids = batch['input_ids'].squeeze(1).to(device) # 将输入数据移动到计算设备上
        attention_mask = batch['attention_mask'].squeeze(1).to(device) # 将输入数据移动到计算设备上
        labels = batch['labels'].to(device) # 将标签移动到计算设备上
        optimizer.zero_grad() # 清空梯度
        embedding = model.bert.embeddings.word_embeddings(input_ids)
        embedding = embedding.detach().clone().requires_grad_(True)
        embedding.retain_grad()  # 保留梯度信息
        outputs = model(inputs_embeds=embedding, attention_mask=attention_mask)
        loss = criterion(outputs.logits, labels)
        embedding_grad = torch.autograd.grad(loss, embedding, allow_unused=True, retain_graph=True)[0]
        perturbed_embedding = fgsm_attack(embedding, epsilon, embedding_grad) # 添加扰动
        perturbed_tokens = tokenizer.convert_ids_to_tokens(np.argmax(perturbed_embedding.detach().cpu().numpy(), axis=-1).tolist()[0])
        perturbed_input_ids = torch.tensor(tokenizer.convert_tokens_to_ids(perturbed_tokens)).unsqueeze(0).to(device) # 将tokens转换为input_ids
        perturbed_outputs = model(perturbed_input_ids, attention_mask=attention_mask, labels=labels) # 模型前向传播
        perturbed_loss = criterion(perturbed_outputs.logits, labels) # 计算对抗样本损失
        optimizer.zero_grad() # 清空梯度
        loss = 0.5 * loss + 0.5 * perturbed_loss # 计算总损失
        loss.backward() # 反向传播，计算总梯度
        optimizer.step() # 更新参数
        pbar(step=step, info={'loss': loss.item()})
        preds = torch.argmax(outputs.logits, dim=1) # 计算原样本预测结果
        train_loss += loss.item() # 累加原样本损失
        train_acc += torch.sum(preds == labels).item() # 计算原样本准确率
    train_loss /= len(train_loader) # 计算平均原样本损失
    train_acc /= len(train_loader.dataset) # 计算平均原样本准确率
    return train_loss, train_acc # 返回训练损失和准确率

def evaluate(model, criterion, test_loader, device):
    """
    测试函数，仅进行前向传播，不生成对抗样本
    :param model: 模型
    :param criterion: 损失函数
    :param test_loader: 测试数据集的数据加载器
    :param device: 计算设备
    :return: 测试损失和准确率
    """
    model.eval() # 设置模型为评估模式
    test_loss = 0
    test_acc = 0
    with torch.no_grad(): # 关闭梯度计算
        for batch in test_loader:
            input_ids = batch['input_ids'].squeeze(1).to(device) # 将输入数据移动到计算设备上
            attention_mask = batch['attention_mask'].squeeze(1).to(device)
            labels = batch['labels'].to(device)
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels) # 模型前向传播
            loss = criterion(outputs.logits, labels) # 计算损失
            test_loss += loss.item() # 加损失
            preds = torch.argmax(outputs.logits, dim=1) # 计算预测结果
            test_acc += torch.sum(preds == labels).item() #计算准确率
    test_loss /= len(test_loader) # 计算平均损失
    test_acc /= len(test_loader.dataset) # 计算平均准确率
    return test_loss, test_acc

5.3 定义测试函数

这是用于测试模型的函数，输入参数包括模型model，损失函数criterion，测试数据集的数据加载器test_loader和计算设备device。该函数会将模型设置为评估模式，关闭梯度计算，进行前向传播，计算损失和准确率，并返回测试损失和准确率。

# 将模型移动到计算设备上（GPU 或 CPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
best_acc = 0 # 初始化最佳准确率为0
for epoch in range(10): # 进行10轮训练
    train_loss, train_acc = train(model, optimizer, criterion, train_loader, device, epsilon=0.3) # 训练模型，并获取训练损失和准确率
    test_loss, test_acc = evaluate(model, criterion, test_loader, device) # 对测试集进行测试，并获取测试损失和准确率
    dev_loss, dev_acc = evaluate(model, criterion, dev_loader, device) # 对验证集进行测试，并获取验证损失和准确率
    print(f'Epoch {epoch+1}, Train Loss {train_loss:.4f}, Train Acc {train_acc:.4f}, Test Loss {test_loss:.4f}, Test Acc {test_acc:.4f}, Dev Loss {dev_loss:.4f}, Dev Acc {dev_acc:.4f}')
    # 打印训练轮数、训练损失和准确率、测试损失和准确率、验证损失和准确率
    if dev_acc > best_acc: # 如果当前验证准确率大于最佳准确率
        best_acc = dev_acc # 更新最佳准确率
        torch.save(model.state_dict(), 'adv_best_model.pt') # 保存模型参数到文件'best_model.pt'

六、Bert+PGD文本分类

FGSM（Fast Gradient Sign Method）和 PGD（Projected Gradient Descent）都是对抗训练中常用的方法，它们的主要区别在于对抗样本的生成方式和训练策略上。

FGSM 是一种基于梯度的对抗样本生成方法，它通过计算损失函数对输入数据的梯度来生成对抗样本。具体来说，对于一个输入样本，FGSM 会计算其梯度，然后将其与一个小的扰动值相乘，从而生成一个对抗样本。FGSM 的优点是计算效率高，但缺点是生成的对抗样本可能不够鲁棒，容易被攻击者攻击。

PGD 是一种基于迭代的对抗样本生成方法，它通过迭代多次生成对抗样本，并在每次迭代中对生成的对抗样本进行投影，以保证其在一定范围内。具体来说，PGD 会在每次迭代中计算输入数据的梯度，然后对其进行一定程度的扰动，并将扰动后的结果进行投影，以保证其在一定范围内。PGD 的优点是生成的对抗样本更加鲁棒，但缺点是计算效率较低。

在 Bert 文本分类任务中，采用 FGSM 和 PGD 对抗训练的区别主要在于训练策略上。FGSM 对抗训练通常采用单次扰动，而 PGD 对抗训练通常采用多次迭代扰动。在 FGSM 对抗训练中，每次训练只使用一个对抗样本，而在 PGD 对抗训练中，每次训练使用多个对抗样本。因此，PGD 对抗训练的鲁棒性更强，但计算代价也更高。

6.1 定义PGD攻击函数

def pgd_attack(model, embedding, attention_mask, labels, epsilon, embedding_grad, alpha, num_iters):
    """
    PGD 攻击函数
    :param model: 模型
    :param embedding: 原始输入的嵌入表示
    :param attention_mask: 输入的注意力掩码
    :param labels: 标签
    :param epsilon: 扰动范围
    :param embedding_grad: 原始输入的嵌入表示的梯度
    :param alpha: 步长
    :param num_iters: 迭代次数
    :return: 添加扰动后的嵌入表示
    """
    perturbed_embedding = torch.nn.Parameter(embedding) # 将嵌入表示转换为可训练的参数
    for i in range(num_iters):
        perturbed_embedding.requires_grad = True # 设置扰动为可求导
        perturbed_outputs = model(inputs_embeds=perturbed_embedding, attention_mask=attention_mask, labels=labels) # 模型前向传播
        perturbed_loss = criterion(perturbed_outputs.logits, labels) # 计算对抗样本损失
        perturbed_grad = torch.autograd.grad(perturbed_loss, perturbed_embedding, allow_unused=True, retain_graph=True)[0] # 计算梯度
        perturbed_embedding = perturbed_embedding.detach() + alpha * torch.sign(perturbed_grad) # 梯度方向上进行一定的步长更新
        perturbed_embedding = torch.max(torch.min(perturbed_embedding, embedding + epsilon), embedding - epsilon) # 将扰动限制在一定范围内
        perturbed_embedding = torch.nn.Parameter(perturbed_embedding) # 将更新后的嵌入表示重新转换为可训练的参数

    return perturbed_embedding

6.2 定义训练函数

def train(model, optimizer, criterion, train_loader, device, epsilon, alpha, num_iters):
    model.train() # 将模型设置为训练模式
    train_loss = 0 # 初始化训练损失为0
    train_acc = 0 # 初始化训练准确率为0
    tokenizer = BertTokenizerFast.from_pretrained('bert-base-chinese')
    for batch in train_loader: # 遍历训练数据集
        input_ids = batch['input_ids'].squeeze(1).to(device) # 将输入数据移动到计算设备上
        attention_mask = batch['attention_mask'].squeeze(1).to(device) # 将输入数据移动到计算设备上
        labels = batch['labels'].to(device) # 将标签移动到计算设备上
        optimizer.zero_grad() # 清空梯度
        embedding = model.bert.embeddings.word_embeddings(input_ids)
        embedding = embedding.detach().clone().requires_grad_(True)
        embedding.retain_grad()  # 保留梯度信息
        outputs = model(inputs_embeds=embedding, attention_mask=attention_mask)
        loss = criterion(outputs.logits, labels)
        embedding_grad = torch.autograd.grad(loss, embedding, allow_unused=True, retain_graph=True)[0]
        perturbed_embedding = pgd_attack(model, embedding, attention_mask, labels, epsilon, embedding_grad, alpha, num_iters)  # 添加扰动
        perturbed_tokens = tokenizer.convert_ids_to_tokens(np.argmax(perturbed_embedding.detach().cpu().numpy(), axis=-1).tolist()[0])
        perturbed_input_ids = torch.tensor(tokenizer.convert_tokens_to_ids(perturbed_tokens)).unsqueeze(0).to(device) # 将tokens转换为input_ids
        perturbed_outputs = model(perturbed_input_ids, attention_mask=attention_mask, labels=labels) # 模型前向传播
        perturbed_loss = criterion(perturbed_outputs.logits, labels) # 计算对抗样本损失
        optimizer.zero_grad() # 清空梯度
        loss = 0.5 * loss + 0.5 * perturbed_loss # 计算总损失
        loss.backward() # 反向传播，计算总梯度
        optimizer.step() # 更新参数
        preds = torch.argmax(outputs.logits, dim=1) # 计算原样本预测结果
        train_loss += loss.item() # 累加原样本损失
        train_acc += torch.sum(preds == labels).item() # 计算原样本准确率
    train_loss /= len(train_loader) # 计算平均原样本损失
    train_acc /= len(train_loader.dataset) # 计算平均原样本准确率
    return train_loss, train_acc # 返回训练损失和准确率

def evaluate(model, criterion, test_loader, device):
    """
    测试函数，仅进行前向传播，不生成对抗样本
    :param model: 模型
    :param criterion: 损失函数
    :param test_loader: 测试数据集的数据加载器
    :param device: 计算设备
    :return: 测试损失和准确率
    """
    model.eval() # 设置模型为评估模式
    test_loss = 0
    test_acc = 0
    with torch.no_grad(): # 关闭梯度计算
        for batch in test_loader:
            input_ids = batch['input_ids'].squeeze(1).to(device) # 将输入数据移动到计算设备上
            attention_mask = batch['attention_mask'].squeeze(1).to(device)
            labels = batch['labels'].to(device)
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels) # 模型前向传播
            loss = criterion(outputs.logits, labels) # 计算损失
            test_loss += loss.item() # 加损失
            preds = torch.argmax(outputs.logits, dim=1) # 计算预测结果
            test_acc += torch.sum(preds == labels).item() #计算准确率
    test_loss /= len(test_loader) # 计算平均损失
    test_acc /= len(test_loader.dataset) # 计算平均准确率
    return test_loss, test_acc

6.3 定义测试函数

# 将模型移动到计算设备上（GPU 或 CPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
best_acc = 0 # 初始化最佳准确率为0
for epoch in range(10): # 进行10轮训练
    train_loss, train_acc = train(model, optimizer, criterion, train_loader, device, epsilon=0.3, alpha=0.04, num_iters=5) # 训练模型，并获取训练损失和准确率
    test_loss, test_acc = evaluate(model, criterion, test_loader, device) # 对测试集进行测试，并获取测试损失和准确率
    dev_loss, dev_acc = evaluate(model, criterion, dev_loader, device) # 对验证集进行测试，并获取验证损失和准确率
    print(f'Epoch {epoch+1}, Train Loss {train_loss:.4f}, Train Acc {train_acc:.4f}, Test Loss {test_loss:.4f}, Test Acc {test_acc:.4f}, Dev Loss {dev_loss:.4f}, Dev Acc {dev_acc:.4f}')
    # 打印训练轮数、训练损失和准确率、测试损失和准确率、验证损失和准确率
    if dev_acc > best_acc: # 如果当前验证准确率大于最佳准确率
        best_acc = dev_acc # 更新最佳准确率
        torch.save(model.state_dict(), 'adv_best_model.pt') # 保存模型参数到文件'best_model.pt'