实战指南：用Python实现PGD对抗攻击提升模型鲁棒性

在图像分类任务中，我们常常会遇到一个令人头疼的问题——精心训练的模型在面对人为设计的微小扰动时，竟然会给出完全错误的预测结果。这种现象被称为对抗样本攻击，而PGD（Projected Gradient Descent）正是生成这类对抗样本的利器之一。本文将带你从零开始，用Python实现一个完整的PGD攻击流程，目标是在CIFAR-10数据集上"欺骗"一个预训练的ResNet模型。

1. 对抗攻击基础与PGD原理

对抗攻击的核心思想是在原始输入上添加人类难以察觉的微小扰动，使得机器学习模型产生错误的输出。PGD作为当前最强大的白盒攻击方法之一，其优势在于：

• 迭代优化 ：通过多轮梯度更新逐步寻找最优扰动
• 投影约束 ：确保扰动始终保持在允许的范围内（如L∞范数约束）
• 强攻击性 ：在相同扰动预算下，通常比单步攻击（如FGSM）更有效

PGD的数学表达可以简化为以下迭代过程：

x_adv = x + delta  # 初始扰动
for i in range(iterations):
    # 1. 计算损失函数关于输入的梯度
    grad = compute_gradient(loss, x_adv) 
    
    # 2. 沿梯度方向更新扰动
    delta = delta + alpha * sign(grad)
    
    # 3. 将扰动投影到允许范围内
    delta = clip(delta, -epsilon, epsilon)
    
    # 4. 确保扰动后的图像仍在有效像素范围内
    x_adv = clip(x + delta, 0, 1)

注意：实际实现时需要考虑图像预处理、模型梯度计算等细节，我们将在代码部分详细展开。

2. 实验环境搭建与数据准备

在开始编写攻击代码前，我们需要准备好实验环境。以下是推荐的环境配置：

# 创建conda环境（可选）
conda create -n pgd_attack python=3.8
conda activate pgd_attack

# 安装核心依赖
pip install torch torchvision numpy matplotlib

对于数据集，我们将使用CIFAR-10和预训练的ResNet-18模型。PyTorch已经内置了这些资源：

import torch
import torchvision
from torchvision import transforms

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
])

# 加载CIFAR-10测试集
testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=32, shuffle=False)

# 加载预训练模型
model = torchvision.models.resnet18(pretrained=True)
model.eval()  # 设置为评估模式

3. PGD攻击的完整实现

现在让我们实现PGD攻击的核心代码。我们将创建一个可配置的PGD攻击类，支持调整关键参数：

class PGDAttack:
    def __init__(self, model, epsilon=8/255, alpha=2/255, 
                 iterations=10, random_start=True):
        self.model = model
        self.epsilon = epsilon  # 扰动上限
        self.alpha = alpha      # 单步扰动大小
        self.iterations = iterations
        self.random_start = random_start
        
    def attack(self, images, labels):
        images = images.clone().detach()
        labels = labels.clone().detach()
        
        # 初始化扰动
        if self.random_start:
            delta = torch.empty_like(images).uniform_(-self.epsilon, self.epsilon)
        else:
            delta = torch.zeros_like(images)
            
        delta = torch.clamp(delta, -self.epsilon, self.epsilon)
        
        # 迭代攻击
        for _ in range(self.iterations):
            delta.requires_grad = True
            
            # 前向传播
            outputs = self.model(images + delta)
            loss = torch.nn.functional.cross_entropy(outputs, labels)
            
            # 反向传播计算梯度
            grad = torch.autograd.grad(loss, delta)[0]
            
            # 更新扰动
            delta = delta.detach() + self.alpha * grad.sign()
            delta = torch.clamp(delta, -self.epsilon, self.epsilon)
            
            # 确保图像仍在有效范围内
            delta = torch.clamp(images + delta, 0, 1) - images
            
        return images + delta

关键参数说明：

4. 攻击效果评估与可视化

实现攻击后，我们需要评估其效果并可视化攻击样本。以下是评估攻击成功率的代码：

def evaluate_attack(model, dataloader, attack):
    correct = 0
    total = 0
    
    for images, labels in dataloader:
        # 生成对抗样本
        adv_images = attack.attack(images, labels)
        
        # 模型预测
        outputs = model(adv_images)
        _, predicted = torch.max(outputs.data, 1)
        
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    accuracy = 100 * correct / total
    print(f'Attack success rate: {100 - accuracy:.2f}%')
    return accuracy

可视化原始图像与对抗样本的对比：

import matplotlib.pyplot as plt

def visualize_attack(original, adversarial, epsilon):
    plt.figure(figsize=(10, 5))
    
    # 原始图像
    plt.subplot(1, 3, 1)
    plt.imshow(original.permute(1, 2, 0))
    plt.title('Original')
    plt.axis('off')
    
    # 对抗样本
    plt.subplot(1, 3, 2)
    plt.imshow(adversarial.permute(1, 2, 0))
    plt.title(f'Adversarial (ε={epsilon})')
    plt.axis('off')
    
    # 扰动放大
    plt.subplot(1, 3, 3)
    perturbation = (adversarial - original).abs().sum(dim=0)
    plt.imshow(perturbation, cmap='hot')
    plt.title('Perturbation Magnified')
    plt.axis('off')
    
    plt.tight_layout()
    plt.show()

5. 防御策略与模型鲁棒性提升

了解攻击方法后，我们自然需要考虑如何防御。以下是几种常见的防御策略：

1. 对抗训练 ：在训练过程中加入对抗样本

   • 优点：直接有效
   • 缺点：训练成本高，可能影响原始准确率

2. 输入预处理 ：

   • 随机化：随机调整大小、填充等
   • 去噪：使用自编码器或滤波技术

3. 梯度掩码 ：

   • 使模型梯度对攻击者不可用或难以利用
   • 但可能只是"虚假的安全感"

以下是简单的对抗训练实现示例：

def adversarial_train(model, train_loader, optimizer, attack, epochs=5):
    model.train()  # 设置为训练模式
    
    for epoch in range(epochs):
        for images, labels in train_loader:
            # 生成对抗样本
            adv_images = attack.attack(images, labels)
            
            # 同时使用原始样本和对抗样本训练
            optimizer.zero_grad()
            outputs = model(torch.cat([images, adv_images]))
            loss = F.cross_entropy(
                outputs, torch.cat([labels, labels]))
            loss.backward()
            optimizer.step()

在实际项目中，我发现对抗训练虽然能显著提升模型鲁棒性，但也需要平衡几个关键因素：

• 训练时间与计算资源的消耗
• 原始测试集准确率的保持
• 对不同攻击方法的泛化能力

一个实用的技巧是从较小的epsilon开始（如4/255），随着训练过程逐步增大，这样可以在保持模型原始性能的同时提升鲁棒性。