对抗AI黑客：大模型安全的终极防御

# 基于对抗训练的大模型防御系统示例
# 运行环境：Python 3.8+, pip install transformers torch numpy pandas scikit-learn matplotlib

import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from transformers import BertForSequenceClassification, BertTokenizer, AdamW, get_linear_schedule_with_warmup
from torch.utils.data import DataLoader, Dataset, TensorDataset
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix
import random
import time
import os

# 设置随机种子，确保结果可复现
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

# 检查是否有GPU可用
def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available():
        return torch.device("mps")
    else:
        return torch.device("cpu")

device = get_device()
print(f"使用设备: {device}")

# 模拟文本分类数据集生成
def generate_text_classification_data(n_samples=1000, n_classes=2):
    """
    生成模拟的文本分类数据集
    """
    # 主题词库
    topics = {
        0: ["技术", "编程", "人工智能", "开发", "算法", "数据", "计算机", "软件", "网络", "安全"],
        1: ["体育", "足球", "篮球", "比赛", "运动员", "冠军", "训练", "教练", "俱乐部", "赛事"],
        2: ["美食", "烹饪", "餐厅", "食材", "食谱", "味道", "健康", "营养", "品尝", "美食节"],
        3: ["旅游", "景点", "旅行", "风景", "酒店", "文化", "体验", "探索", "冒险", "度假"]
    }
    
    # 确保类别数量不超过主题数量
    n_classes = min(n_classes, len(topics))
    
    # 生成数据
    texts = []
    labels = []
    
    for i in range(n_samples):
        # 随机选择标签
        label = random.randint(0, n_classes - 1)
        
        # 生成文本：3-6个主题词，加上随机填充词
        n_topic_words = random.randint(3, 6)
        topic_words = random.sample(topics[label], min(n_topic_words, len(topics[label])))
        
        # 添加随机填充词
        filler_words = ["这是", "一个", "关于", "非常", "重要", "的", "有", "很多", "人们", "喜欢"]
        n_filler_words = random.randint(2, 5)
        selected_filler_words = random.sample(filler_words, n_filler_words)
        
        # 混合所有词并打乱顺序
        all_words = topic_words + selected_filler_words
        random.shuffle(all_words)
        
        # 组成句子
        text = """".join(all_words[:3]) + "是" + """".join(all_words[3:]) + "。"  # 简单的句子结构
        
        texts.append(text)
        labels.append(label)
    
    return pd.DataFrame({"text": texts, "label": labels})

# 加载预训练模型和分词器
def load_model(model_name="bert-base-chinese", num_labels=2):
    """
    加载预训练的BERT模型和分词器
    """
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
    model.to(device)
    return model, tokenizer

# 数据预处理和加载
def preprocess_data(data, tokenizer, max_length=128):
    """
    预处理数据并创建DataLoader
    """
    # 分词
    inputs = tokenizer(
        data["text"].tolist(),
        max_length=max_length,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    
    # 创建数据集
    labels = torch.tensor(data["label"].tolist())
    dataset = TensorDataset(inputs["input_ids"], inputs["attention_mask"], labels)
    
    return dataset

# FGSM攻击实现
def fgsm_attack(model, inputs, labels, epsilon=0.01):
    """
    使用FGSM算法生成对抗样本
    """
    # 将输入和标签移至设备
    input_ids = inputs["input_ids"].to(device)
    attention_mask = inputs["attention_mask"].to(device)
    labels = labels.to(device)
    
    # 确保可以计算梯度
    input_ids.requires_grad = True
    
    # 前向传播
    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
    loss = outputs.loss
    
    # 反向传播计算梯度
    model.zero_grad()
    loss.backward()
    
    # 获取输入的梯度
    input_grad = input_ids.grad.data
    
    # 根据梯度符号生成对抗扰动
    perturbed_input = input_ids + epsilon * input_grad.sign()
    
    # 确保扰动后的输入仍然是有效的token ID（这是简化版，实际应用中需要更复杂的处理）
    perturbed_input = torch.clamp(perturbed_input, min=0, max=tokenizer.vocab_size - 1).type(torch.long)
    
    return perturbed_input, attention_mask, labels

# PGD攻击实现
def pgd_attack(model, inputs, labels, epsilon=0.01, alpha=0.001, steps=40):
    """
    使用PGD算法生成对抗样本
    """
    # 将输入和标签移至设备
    input_ids = inputs["input_ids"].to(device).detach()
    attention_mask = inputs["attention_mask"].to(device)
    labels = labels.to(device)
    
    # 初始化扰动
    perturbed_input = input_ids.clone().detach()
    perturbed_input.requires_grad = True
    
    # 在多个步骤中迭代生成对抗样本
    for _ in range(steps):
        # 前向传播
        outputs = model(input_ids=perturbed_input, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        
        # 反向传播计算梯度
        model.zero_grad()
        loss.backward()
        
        # 获取输入的梯度
        input_grad = perturbed_input.grad.data
        
        # 根据梯度方向更新扰动
        perturbed_input = perturbed_input + alpha * input_grad.sign()
        
        # 裁剪扰动范围
        delta = torch.clamp(perturbed_input - input_ids, min=-epsilon, max=epsilon)
        perturbed_input = torch.clamp(input_ids + delta, min=0, max=tokenizer.vocab_size - 1).type(torch.long).detach()
        perturbed_input.requires_grad = True
    
    return perturbed_input, attention_mask, labels

# 对抗训练函数
def adversarial_training(model, train_loader, val_loader, tokenizer, optimizer, scheduler, num_epochs=5, attack_method="fgsm", epsilon=0.01):
    """
    执行对抗训练
    """
    # 记录训练过程
    train_losses = []
    val_losses = []
    train_accs = []
    val_accs = []
    
    # 选择攻击方法
    if attack_method == "fgsm":
        attack_fn = fgsm_attack
    elif attack_method == "pgd":
        attack_fn = pgd_attack
    else:
        raise ValueError(f"不支持的攻击方法: {attack_method}")
    
    # 开始训练
    for epoch in range(num_epochs):
        print(f"\nEpoch {epoch+1}/{num_epochs}")
        print("=" * 50)
        
        # 训练模式
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        start_time = time.time()
        
        for batch_idx, (input_ids, attention_mask, labels) in enumerate(train_loader):
            # 准备输入
            inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
            
            # 生成对抗样本
            perturbed_input, perturbed_mask, perturbed_labels = attack_fn(model, inputs, labels, epsilon=epsilon)
            
            # 合并原始样本和对抗样本
            combined_inputs = torch.cat([input_ids.to(device), perturbed_input], dim=0)
            combined_masks = torch.cat([attention_mask.to(device), perturbed_mask], dim=0)
            combined_labels = torch.cat([labels.to(device), perturbed_labels], dim=0)
            
            # 前向传播
            outputs = model(input_ids=combined_inputs, attention_mask=combined_masks, labels=combined_labels)
            loss = outputs.loss
            logits = outputs.logits
            
            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()
            
            # 统计损失和准确率
            train_loss += loss.item() * combined_inputs.size(0)
            _, predicted = torch.max(logits, 1)
            train_correct += (predicted == combined_labels).sum().item()
            train_total += combined_labels.size(0)
            
            # 打印进度
            if (batch_idx + 1) % 10 == 0:
                print(f"Batch {batch_idx+1}/{len(train_loader)}, Train Loss: {loss.item():.4f}")
        
        # 计算平均训练损失和准确率
        train_loss /= train_total
        train_acc = train_correct / train_total
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        
        # 验证模式
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for input_ids, attention_mask, labels in val_loader:
                # 移至设备
                input_ids = input_ids.to(device)
                attention_mask = attention_mask.to(device)
                labels = labels.to(device)
                
                # 前向传播
                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                loss = outputs.loss
                logits = outputs.logits
                
                # 统计损失和准确率
                val_loss += loss.item() * input_ids.size(0)
                _, predicted = torch.max(logits, 1)
                val_correct += (predicted == labels).sum().item()
                val_total += labels.size(0)
        
        # 计算平均验证损失和准确率
        val_loss /= val_total
        val_acc = val_correct / val_total
        val_losses.append(val_loss)
        val_accs.append(val_acc)
        
        end_time = time.time()
        
        print(f"Epoch {epoch+1} Summary:")
        print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
        print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")
        print(f"Time taken: {end_time - start_time:.2f} seconds")
    
    return model, {
        "train_losses": train_losses,
        "val_losses": val_losses,
        "train_accs": train_accs,
        "val_accs": val_accs
    }

# 评估模型在对抗样本上的性能
def evaluate_adversarial_robustness(model, test_loader, tokenizer, attack_method="fgsm", epsilon=0.01):
    """
    评估模型在对抗样本上的鲁棒性
    """
    # 选择攻击方法
    if attack_method == "fgsm":
        attack_fn = fgsm_attack
    elif attack_method == "pgd":
        attack_fn = pgd_attack
    else:
        raise ValueError(f"不支持的攻击方法: {attack_method}")
    
    # 评估模式
    model.eval()
    
    # 记录结果
    original_correct = 0
    adversarial_correct = 0
    total = 0
    
    # 存储原始预测和对抗预测
    original_predictions = []
    adversarial_predictions = []
    true_labels = []
    
    with torch.no_grad():
        for input_ids, attention_mask, labels in test_loader:
            # 准备输入
            inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
            
            # 移至设备
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            labels = labels.to(device)
            
            # 原始样本的预测
            original_outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            _, original_pred = torch.max(original_outputs.logits, 1)
            
            # 生成对抗样本
            perturbed_input, perturbed_mask, _ = attack_fn(model, inputs, labels, epsilon=epsilon)
            
            # 对抗样本的预测
            adversarial_outputs = model(input_ids=perturbed_input, attention_mask=perturbed_mask)
            _, adversarial_pred = torch.max(adversarial_outputs.logits, 1)
            
            # 统计正确数
            original_correct += (original_pred == labels).sum().item()
            adversarial_correct += (adversarial_pred == labels).sum().item()
            total += labels.size(0)
            
            # 存储预测结果
            original_predictions.extend(original_pred.cpu().numpy())
            adversarial_predictions.extend(adversarial_pred.cpu().numpy())
            true_labels.extend(labels.cpu().numpy())
    
    # 计算准确率
    original_acc = original_correct / total
    adversarial_acc = adversarial_correct / total
    
    # 计算攻击成功率
    attack_success_rate = 1 - adversarial_acc
    
    print(f"原始样本准确率: {original_acc:.4f}")
    print(f"对抗样本准确率: {adversarial_acc:.4f}")
    print(f"攻击成功率: {attack_success_rate:.4f}")
    
    # 计算F1分数
    original_f1 = f1_score(true_labels, original_predictions, average="weighted")
    adversarial_f1 = f1_score(true_labels, adversarial_predictions, average="weighted")
    
    print(f"原始样本F1分数: {original_f1:.4f}")
    print(f"对抗样本F1分数: {adversarial_f1:.4f}")
    
    # 生成混淆矩阵
    original_cm = confusion_matrix(true_labels, original_predictions)
    adversarial_cm = confusion_matrix(true_labels, adversarial_predictions)
    
    return {
        "original_acc": original_acc,
        "adversarial_acc": adversarial_acc,
        "attack_success_rate": attack_success_rate,
        "original_f1": original_f1,
        "adversarial_f1": adversarial_f1,
        "original_cm": original_cm,
        "adversarial_cm": adversarial_cm
    }

# 可视化训练过程

def visualize_training_history(history):
    """
    可视化训练过程中的损失和准确率
    """
    epochs = range(1, len(history["train_losses"]) + 1)
    
    # 创建图形
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # 绘制损失曲线
    ax1.plot(epochs, history["train_losses"], 'b-o', label="Training Loss")
    ax1.plot(epochs, history["val_losses"], 'r-s', label="Validation Loss")
    ax1.set_title("Training and Validation Loss")
    ax1.set_xlabel("Epochs")
    ax1.set_ylabel("Loss")
    ax1.legend()
    
    # 绘制准确率曲线
    ax2.plot(epochs, history["train_accs"], 'b-o', label="Training Accuracy")
    ax2.plot(epochs, history["val_accs"], 'r-s', label="Validation Accuracy")
    ax2.set_title("Training and Validation Accuracy")
    ax2.set_xlabel("Epochs")
    ax2.set_ylabel("Accuracy")
    ax2.legend()
    
    plt.tight_layout()
    
    # 在实际应用中，可以保存图像或显示
    # plt.savefig('training_history.png')
    # plt.show()
    
    return plt

# 主函数
def main():
    print("开始大模型对抗训练防御系统演示...")
    
    # 1. 生成模拟数据
    print("\n1. 生成模拟文本分类数据集...")
    data = generate_text_classification_data(n_samples=2000, n_classes=2)
    print(f"生成的数据量: {len(data)}")
    print("数据样例:")
    print(data.head())
    
    # 2. 划分训练集、验证集和测试集
    print("\n2. 划分数据集...")
    train_data, test_data = train_test_split(data, test_size=0.2, random_state=42)
    train_data, val_data = train_test_split(train_data, test_size=0.2, random_state=42)
    print(f"训练集大小: {len(train_data)}")
    print(f"验证集大小: {len(val_data)}")
    print(f"测试集大小: {len(test_data)}")
    
    # 3. 加载预训练模型和分词器
    print("\n3. 加载预训练模型和分词器...")
    model, tokenizer = load_model(model_name="bert-base-chinese", num_labels=2)
    
    # 4. 数据预处理
    print("\n4. 预处理数据...")
    train_dataset = preprocess_data(train_data, tokenizer)
    val_dataset = preprocess_data(val_data, tokenizer)
    test_dataset = preprocess_data(test_data, tokenizer)
    
    # 创建DataLoader
    batch_size = 16
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    # 5. 设置优化器和学习率调度器
    print("\n5. 设置优化器和学习率调度器...")
    optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)
    num_epochs = 3
    total_steps = len(train_loader) * num_epochs
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps)
    
    # 6. 执行对抗训练
    print("\n6. 执行对抗训练...")
    model, history = adversarial_training(
        model, 
        train_loader, 
        val_loader, 
        tokenizer, 
        optimizer, 
        scheduler, 
        num_epochs=num_epochs,
        attack_method="fgsm",
        epsilon=0.01
    )
    
    # 7. 可视化训练历史
    print("\n7. 可视化训练历史...")
    plt = visualize_training_history(history)
    print("训练历史可视化完成")
    
    # 8. 评估模型的对抗鲁棒性
    print("\n8. 评估模型的对抗鲁棒性...")
    # 评估FGSM攻击下的鲁棒性
    print("\n评估FGSM攻击下的鲁棒性:")
    fgsm_results = evaluate_adversarial_robustness(model, test_loader, tokenizer, attack_method="fgsm", epsilon=0.01)
    
    # 评估PGD攻击下的鲁棒性
    print("\n评估PGD攻击下的鲁棒性:")
    pgd_results = evaluate_adversarial_robustness(model, test_loader, tokenizer, attack_method="pgd", epsilon=0.01)
    
    # 9. 保存模型
    print("\n9. 保存训练好的模型...")
    model_save_path = "adversarially_trained_model"
    if not os.path.exists(model_save_path):
        os.makedirs(model_save_path)
    
    # 保存模型和分词器
    model.save_pretrained(model_save_path)
    tokenizer.save_pretrained(model_save_path)
    print(f"模型已保存到: {model_save_path}")
    
    print("\n大模型对抗训练防御系统演示完成！")

if __name__ == "__main__":
    main()