pytorch框架实现BIM算法

最新推荐文章于 2024-07-16 22:49:14 发布

weixin_57686258

最新推荐文章于 2024-07-16 22:49:14 发布

阅读量186

点赞数

文章标签： pytorch 算法深度学习

本文链接：https://blog.csdn.net/weixin_57686258/article/details/130786165

版权

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

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

# 定义模型
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

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()

    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

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