NKCTF有两道ai题，遗憾的是比赛时还并未涉足ai安全。发现第一题是很经典DLG梯度泄露故尝试复现，毕竟出题人连出题代码都给了，不复现一下怎么行呢。
一开始我先去了解了一下ai安全，知识获取渠道一个是alice佬
另一个是安全客上一位大牛
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佬的博客深耕一下这个方向
