MI-FGSM代码及注释

import os
from PIL import Image
import torch
from torch import nn
import torchvision.transforms as transforms

def mi_fgsm_attack(image, epsilon, data_grad, alpha, decay_factor, g):
    # 计算梯度的符号
    sign_data_grad = data_grad.sign()
    # 更新累积梯度
    g = decay_factor * g + sign_data_grad
    # 通过调整输入图像的每个像素来创建扰动图像
    perturbed_image = image + alpha * g
    # 将扰动图像剪切到[0,1]范围内
    perturbed_image = torch.clamp(perturbed_image, 0, 1)
    # 添加扰动
    perturbed_image = torch.clamp(perturbed_image, image - epsilon, image + epsilon)
    # 返回扰动图像和累积梯度
    return perturbed_image, g

# 定义模型
model = nn.Sequential(
    nn.Conv2d(1, 32, kernel_size=3),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(32, 64, kernel_size=3),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Flatten(),
    nn.Linear(1600, 10)
)

# 加载预训练模型权重
model.load_state_dict(torch.load('model.pth'))

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 定义图像转换
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor()
])

# 定义图像文件夹路径
image_folder = 'images'

# 获取图像文件列表
image_files = os.listdir(image_folder)

# 迭代次数
num_iter = 10

# 迭代步长
alpha = 0.05

# 最大扰动值
epsilon = 0.1

# 衰减因子
decay_factor = 1.0

for image_file in image_files:
    # 获取图像文件路径
    image_path = os.path.join(image_folder, image_file)

    # 读取图像文件并转换为灰度图像
    image = Image.open(image_path).convert('L')

    # 转换图像并添加批次维度
    image = transform(image).unsqueeze(0)

    # 将模型设置为评估模式
    model.eval()

    # 定义扰动图像和累积梯度
    perturbed_image = image.clone().detach()
    g = torch.zeros_like(perturbed_image)

    for i in range(num_iter):
        # 前向传播以获取预测结果
        output = model(perturbed_image)

        # 获取预测标签
        label = output.argmax(dim=1)

        # 计算损失
        loss = criterion(output, label)

        # 反向传播以获取梯度
        model.zero_grad()
        loss.backward()
        data_grad = perturbed_image.grad.data

        # 调用 mi_fgsm_attack 函数生成对抗样本和更新累积梯度
        perturbed_image, g = mi_fgsm_attack(perturbed_image, epsilon, data_grad, alpha, decay_factor, g)