AI学习之路(2)–一次半失败的安全ai题复现

NKCTF有两道ai题，遗憾的是比赛时还并未涉足ai安全。发现第一题是很经典DLG梯度泄露故尝试复现，毕竟出题人连出题代码都给了，不复现一下怎么行呢。

一开始我先去了解了一下ai安全，知识获取渠道一个是alice佬

Alice’s challenge- CSDN搜索

另一个是安全客上一位大牛

https://www.anquanke.com/member.html?memberId=157437

学习了相关知识后尝试复现，wp中的代码并不能copy即用，大概有两个问题一是pillow版本导致其中一个函数无法使用，另一个是需要自行配置一下LeNet（我是直接拿laogong队wp里的LeNet）。

修改代码如下

from PIL import Image, ImageDraw, ImageFont
characters = "0123456789qwertyuiopasdfghjklzxcvbnm"
image_size = (32, 32)
for idx,character in enumerate(characters):
 image = Image.new("RGB", image_size, (0,0,0))
 draw = ImageDraw.Draw(image)
 font = ImageFont.truetype(r"C:\Windows\Fonts\Inkfree.ttf", 30)
 left, top, right, bottom = draw.textbbox((0, 0), character, font)
 width, height = right - left, bottom - top
 x = (image_size[0] - width) // 2
 y = (image_size[1] - height) // 2
 draw.text((x, y), character, font=font, fill=(255, 255, 255), stroke_width=1, stroke_fill=(255, 255, 255))
 image.save(fr"data\{idx}.png")

# 导入所需模块
import torch
from PIL import Image
import os
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import pickle
# 定义标签转为one-hot编码函数
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        act = nn.Sigmoid
        self.body = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=5, padding=5 // 2, stride=2),
            act(),
            nn.Conv2d(32, 32, kernel_size=5, padding=5 // 2, stride=2),
            act(),
            nn.Conv2d(32, 64, kernel_size=5, padding=5 // 2, stride=1),
            act(),
        )
        self.fc = nn.Sequential(
            nn.Linear(4096, 36)
        )

    def forward(self, x):
        x = self.body(x)
        x = x.flatten(0).unsqueeze(0)
        x = self.fc(x)
        return x
def label_to_onehot(target):
    target = torch.unsqueeze(target, 1)  # 增加一个维度，便于后面矩阵操作
    onehot_target = torch.zeros(target.size(0), 36, device=target.device)  # 初始化大小为目标大小的全零张量
    onehot_target.scatter_(1, target, 1)  # 使用scatter_方法生成one-hot编码
    return onehot_target
# 定义针对one-hot编码标签的交叉熵损失函数
def cross_entropy_for_onehot(pred, target):
    return torch.mean(torch.sum(-target * F.log_softmax(pred, dim=-1), 1))  # 计算平均交叉熵损失
# 初始化模型权重函数
def weights_init(m):
    if hasattr(m, "weight"):  # 如果层有weight属性
        m.weight.data.uniform_(-0.5, 0.5)  # 随机初始化权重
    if hasattr(m, "bias"):  # 如果层有bias属性
        m.bias.data.uniform_(-0.5, 0.5)  # 随机初始化偏置
# 主函数入口
if __name__ == "__main__":
    # 定义一个字符串用于后续验证
    flag = "f1ag1scainpeace114514"
    assert len(flag) == 21
    # 设定设备（这里使用CPU）
    device = "cpu"
    # 设置损失函数为自定义的one-hot交叉熵损失函数
    criterion = cross_entropy_for_onehot
    # 定义字符字典
    dic = "0123456789qwertyuiopasdfghjklzxcvbnm"
    # 初始化图像转换器
    tt = transforms.ToTensor()  # 图片转Tensor
    tp = transforms.ToPILImage()  # Tensor转图片
    # 遍历字典中的每个字符
    for i, c in enumerate(dic):
        # 设置随机种子确保每次实验结果可复现
        torch.manual_seed(1234)
        # 创建并初始化LeNet模型
        model = LeNet()
        model.to(device)
        model.apply(weights_init)
        # 获取当前字符在字典中的索引
        idx = dic.index(c)
        # 打开对应索引的图像文件
        img = Image.open(f'data/{idx}.png')
        # img = img.convert("L")  # 若原图不是灰度图，可以使用这行代码将其转为灰度图
        # 将图片转换为Tensor格式
        img_tensor = tt(img)
        # 创建目标标签，并将其转换为LongTensor类型，同时移动到指定设备上
        gt_label = torch.Tensor([int(idx)]).long().to(device)
        gt_label = gt_label.view(1, )  # 调整形状为(1,)
        # 将目标标签转换为one-hot编码
        gt_onehot_label = label_to_onehot(gt_label)
        # 输入图片至模型计算预测输出
        out = model(img_tensor)
        # 计算损失
        y = criterion(out, gt_onehot_label)
        # 计算损失相对于模型参数的梯度
        dy_dx = torch.autograd.grad(y, model.parameters())
        # 保存原始梯度信息
        original_dy_dx = list((_.detach().clone() for _ in dy_dx))
        # 将原始梯度信息序列化并保存为.pkl文件
        with open(f'{i}.pkl', 'wb') as f:
            pickle.dump(original_dy_dx, f)

至此复现出一半，由于缺少附件故复现到此为止。整个过程还是能学到很多东西的，后面希望可以跟着Alice佬的博客深耕一下这个方向