从0学习卷积神经网络---强网杯EZNN writeup
2023-12-26 14:59:56 Author: xz.aliyun.com(查看原文) 阅读量:15 收藏

背景

今年强网杯出了一道关于神经网络的Pwn题,当时看到感觉一脸懵逼。赛后看了writeup感觉也是一知半解。感觉个人在算法和模型上的知识欠缺较多,遂借这个机会学习下神经网络的相关知识。

神经网络介绍

神经网络是一种模拟人脑神经元工作原理的计算模型,它由大量的人工神经元相互连接构成,能够学习和处理各种复杂的数据模式。

神经网络的基本结构

  • 神经网络由多个层组成,包括输入层、隐藏层和输出层。
  • 输入层接收原始数据,隐藏层进行中间计算和特征提取,输出层生成最终的预测或分类结果。
  • 神经元是神经网络的基本组成单元,它们在各层之间形成复杂的连接结构。

神经元的结构和功能

  • 每个神经元接收来自前一层神经元的多个输入信号,并对这些信号进行处理以产生一个输出信号。
  • 神经元的输入和输出都是实数值,可以代表各种类型的数据,如图像像素、文本特征等。

神经元的计算过程

  • 加权求和阶段

    • 对于每个输入信号,神经元都有一个对应的权重值(weight)。这些权重决定了输入信号对当前神经元输出的影响程度。
    • 神经元将所有输入信号与它们的权重相乘,然后将这些乘积相加。这个加权求和的结果被称为“净输入”(net input)或“线性组合”(linear combination)。
    net_input = Σ(input_i * weight_i) + bias

    其中,input_i 是第 i 个输入信号,weight_i 是对应输入信号的权重,bias 是一个偏置项,用于调整神经元的激活门槛。

  • 激活函数阶段

    • 加权求和的结果通过一个非线性函数(激活函数)进行变换,生成神经元的输出信号。
    • 激活函数的作用是引入非线性特性到神经网络中,使得网络能够学习和表示更复杂的函数关系。
    output = activation_function(net_input)

    其中,activation_function 是所选的非线性函数,常见的激活函数包括sigmoid、ReLU(Rectified Linear Unit)、tanh、softmax等。

神经网络的前向传播和反向传播

  • 在前向传播过程中,输入数据从输入层经过隐藏层(可能有多层)到达输出层,每个神经元都执行上述的加权求和和激活函数计算,生成网络的预测结果。
  • 在反向传播过程中,首先计算预测结果与实际目标之间的误差,然后这个误差信息通过梯度下降或其他优化算法反向传播到网络中的每个神经元,用于更新权重参数,以减小预测误差和改进模型性能。

神经网络通过这种神经元的计算过程以及前向传播和反向传播的学习机制,能够在各种任务中表现出强大的学习和泛化能力,包括图像识别、语音识别、自然语言处理、推荐系统等领域。

卷积神经网络介绍

卷积神经网络(Convolutional Neural Networks, CNN)是一种专门设计用于处理图像、视频、语音等具有网格结构数据的深度学习模型。

主要组件

  • 卷积层:这是CNN的核心层,使用一组可学习的过滤器(或称卷积核)在输入数据上进行滑动窗口操作,并执行元素级别的乘法和加法运算,生成特征图(或称为激活图)。这些特征图捕获输入中的局部特征和模式。

  • 激活函数:如ReLU(Rectified Linear Unit)、sigmoid或tanh,用于引入非线性特性,使得网络能够学习更复杂的模式。

  • 池化层:通常在卷积层之后,池化层通过取每个子区域的最大值、平均值或其他统计量来降低数据维度和减少计算复杂性,同时保留重要的特征信息。

  • 全连接层:经过一系列的卷积和池化层后,可能会添加一个或多个全连接层。这些层中,每个神经元与前一层的所有神经元都有连接,用于将学到的局部特征组合成全局表示,并用于分类或回归任务。

优势

  • 参数共享:在同一特征图中,卷积核在所有位置重复使用,减少了模型的参数数量,提高了计算效率并增强了模型的泛化能力。

  • 平移不变性:由于卷积操作对输入数据的位置相对不敏感,CNN能够在一定程度上保持对对象位置变化的不变性。

  • 层次化表示:随着网络深度的增加,CNN能够学习越来越抽象和复杂的特征,从边缘和纹理到高级物体部分和整体形状。

应用领域

  • 图像分类和识别
  • 物体检测和分割
  • 人脸识别和行人检测
  • 语音识别和音频处理
  • 自然语言处理中的词嵌入和文本分类

训练过程

  • 使用反向传播算法计算损失函数关于每个参数的梯度。
  • 利用优化算法(如梯度下降、随机梯度下降或Adam)更新网络的权重和偏置,以最小化损失函数。
  • 训练过程通常包括多个迭代周期(epochs),直到模型在验证集上的性能不再显著提高或者达到预设的停止条件。

卷积神经网络因其在处理图像和其他网格结构数据方面的优秀性能和高效性而被广泛应用于各种计算机视觉和深度学习任务中。

卷积神经网络模型计算过程

卷积神经网络(CNN)的计算过程可以分为以下几个主要步骤:

输入预处理

  • 图像通常需要进行一些预处理,如归一化(将像素值调整到0-1之间或-1到1之间)、尺寸调整等,以便于网络处理。

卷积层计算

  • 在卷积层中,每个神经元使用一个可学习的 滤波器(卷积核) 与输入图像的一部分(感受野)进行卷积操作。
  • 卷积操作包括元素级别的乘法和加法运算。将卷积核与输入图像的感受野对应元素相乘,然后将乘积累加起来得到一个单一的输出值。
  • 卷积核在输入图像上以一定的步长滑动,每次滑动后执行一次卷积运算。
  • 每个卷积核会生成一个特征图,反映了输入数据中特定特征的存在和强度。
  • 一个卷积层可能包含多个卷积核,每个卷积核产生一个特征图,因此输出是一组特征图。

激活函数应用

  • 在每个卷积层的输出上,通常会应用一个非线性激活函数,如ReLU、sigmoid或tanh。
  • 激活函数引入了非线性特性,使得网络能够学习更复杂的模式。

池化层计算

  • 池化层用于降低数据维度和减少计算复杂性,同时保留重要的特征信息。
  • 最常用的池化操作是最大池化和平均池化,它们分别取每个子区域的最大值和平均值作为输出。
  • 池化窗口在特征图上以一定的步长滑动,并对每个位置执行池化操作。

重复卷积和池化层

  • 上述卷积、激活和池化过程可能会重复多次,形成多层的卷积结构。每一层都会提取更复杂、更抽象的特征。

平坦化和全连接层

  • 在一系列卷积和池化层之后,通常会将特征图展平成一维向量,然后将其输入到全连接层。
  • 全连接层中的每个神经元与前一层的所有神经元都有连接,这些层主要用于分类或回归任务。

输出层

  • 输出层的结构取决于具体任务。对于分类问题,可能是一个softmax层,输出每个类别的概率分布;对于回归问题,可能是一个线性输出层。

前向传播

  • 前向传播是从输入数据通过网络计算输出的过程。在每个阶段,都按照上述步骤进行计算。

损失函数计算

  • 计算模型的预测输出与实际标签之间的差异,通常使用交叉熵损失函数(对于分类问题)或其他适合特定任务的损失函数。

反向传播和优化

  • 使用反向传播算法计算损失函数关于每个参数的梯度。
  • 然后,使用优化算法(如梯度下降、随机梯度下降或Adam)更新网络的权重和偏置,以最小化损失函数。

重复训练

  • 这个过程(前向传播、反向传播和权重更新)会反复进行许多次(epochs),直到网络的性能在验证集上不再显著提高或者达到预设的停止条件。

通过以上步骤,卷积神经网络能够自动从输入数据中学习和提取有用的特征,并用于各种任务,如图像分类、物体检测、语义分割和语音识别等。

题目分析

第一遍分析题目,可以发现题目的验证流程大概如下:

  1. 开始验证了一个proof-of-work
  2. 读取了一些输入
  3. 以输入建立的模型,对测试集(./example/%lu_%lu)进行测试。在经历30次测试,且测试结果与预期结果匹配了25次以上时,进入下一轮

以上部分是第一遍分析题目时可以得到的最基本的信息。再加上题目提示 "hint for you, norm, 0.1307, 0.3081" , 检索0.1307,0.3081这两个关键浮点数会发现这两个浮点数是MNIST数据集标准化的常用参数。

  • MNIST数据集是机器学习领域中非常经典的一个数据集,由60000个训练样本和10000个测试样本组成,每个样本都是一张28 * 28像素的灰度手写数字图片,其中每一张图片都代表0~9中的一个数字。

得到以上信息之后,结合题目名 EZNN(Neural Networks) 可以看出,题目的前半部分应该是要求输入一个神经网络的一些参数,然后以这些参数构建模型去运行对应的测试集。问题的关键在于如何正确的对程序进行输入,这里就需要对程序进行逆向。

根据字符串的一些提示,我们可以大概推断以下几个函数的名称或作用:

  • sub_11ECC0:运行模型,根据输入图片得到对应的预测值(根据函数调用后的输出字符串可得)
  • sub_121670: doConv,做卷积计算(根据函数内字符串提示可得)
  • sub_122CA0: doPoing,进行池化(根据函数内字符串提示可得)

这里我们可以先尝试用pytroch提供的针对MINIST的数据集的模型样例跑一下,尝试训练我们第一个卷积神经网络的模型(安装完几个依赖库可以直接跑起来):


得到第一个模型之后我们需要确认,我们如何将这个模型输入给EZNN这个程序。这时候可以看一下pytorch提供的导出模型的接口,具体导出了哪些数据。导出一个CNN模型的样例:

torch.save(model.state_dict(), "mnist_cnn.pt")

可以看到主要保存的是模型的 state_dict(), state_dict() 里有什么呢?我们可以打印看一下:

import torch

state_dict = torch.load('mnist_cnn.pt')
print(state_dict)

看一下对应的输出,你会发现它包含了conv1.weight,conv1.bias,conv2.bias,fc1.weight,fc1.bias,fc2.weight,fc2.bias。参考上面卷积神经网络的计算过程以及模型样例的代码,可以确认,这些就是最终训练得到的模型关键参数。

这时候可以对比一下模型样例中forward的实现:

def forward(self, x):
        x = self.conv1(x) #进行第一层卷积层的计算
        x = F.relu(x)     #激活函数
        x = self.conv2(x) #进行第二层卷积层的计算
        x = F.relu(x)     #激活函数
        x = F.max_pool2d(x, 2) #池化,kernel size是2
        x = self.dropout1(x) #丢弃
        x = torch.flatten(x, 1) #在进行连接层之前进行扁平化
        x = self.fc1(x) #连接层第一次计算
        x = F.relu(x)   #激活函数
        x = self.dropout2(x) #丢弃
        x = self.fc2(x) #连接层第二次计算
        output = F.log_softmax(x, dim=1)#得到最后的输出 
        return output
  • 注:以上函数实现的数学细节这里不过多展开,我只会罗列一些解题过程中需要涉及的概念;对具体计算过程及公式感兴趣的朋友可以自行搜索

和题目给出的sub_11ECC0(运行模型)的函数进行一个简单对比。可以粗略看出两者十分类似,都是对输入进行一系列的卷积,池化操作,最后进行全连接层(FC Layer)的计算。题目的输入前读取的两个整形变量最终分别会决定在sub_11ECC0(运行模型)中进行卷积计算的次数和全连接层计算的次数。这里大致可以猜出,我们的输入应该就是pytorch导出的state_dict中的所有数据了。

直接发送模型样例代码训练出来的模型会发现输出“invalid!”。这时候需要对比下第一部分卷积层计算的部分样例代码和题目环境的不同,会发现卷积层使用的卷积核大小及参数不同:

def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)#参数分别为:in_channels,out_channels,kernel_size,stride
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

而题目中生成卷积核的函数为sub_123F60,通过逆向及参数分析可以看出题目创建的所有卷积核参数均为:输入通道为1,输出通道为1,3*3大小。结合程序逆向,我们修改样例模型为:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 1, 3, 1)
        self.conv2 = nn.Conv2d(1, 1, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(121, 121)
        self.fc2 = nn.Linear(121, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        x = F.relu(x)
        output = F.log_softmax(x, dim=1)
        return output

以该模型进行输入之后,可以看到已经可以正常过第一次模型的check了:


这是可以发现题目输出了十个bacodoor的数据,标签,及其预测的标签。之后又是一组类似于第一步的输入。很明显,题目要求我们输入一个新模型,该新模型在题目给出的十个backdoor图片上应该符合题目要求的pridict label(即将图片识别为一个错误的结果)。同时我们通过逆向分析可以看到最后还会使用我们输入的第二个模型去跑目录在./testset/下的数据集。至此我们可以总结一下题目的大致流程:

  • 输入第一个模型,程序使用一些数据集进行测试
  • 程序给出backdoor的数据及对应的标签
  • 输入第二个模型,程序使用backdoor数据集进行测试
  • 程序使用正常数据集进行测试
  • 测试均通过,打印flag

至此,这个题目的目的就是考察选手训练一个能识别出后门数据的模型。因为在backdoor数据集之后还会测试正常数据,这里我们可以考虑直接增大backdoor数据集在总训练集中的比例(或用backdoor多训几轮),来让该模型可以记住对应的backdoor。

剩下的一些工作就集中在数据处理等部分了,这里就不详细描述了,exp最终完成如下:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
import hashlib
from pwn import *

#context.log_level = 'debug'

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 1, 3, 1)
        self.conv2 = nn.Conv2d(1, 1, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(121, 121)
        self.fc2 = nn.Linear(121, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        x = F.relu(x)
        output = F.log_softmax(x, dim=1)
        return output

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

#用backdoor的数据结合minst数据集进行训练
def train_with_backdoor(data):
    torch.manual_seed(1)

    device = torch.device("cpu")

    train_kwargs = {'batch_size': 64,'shuffle': True}
    test_kwargs = {'batch_size': 1000}
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    mixed = torch.utils.data.ConcatDataset([dataset1, data * 2000])
    train_loader = torch.utils.data.DataLoader(mixed,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
    backdoor_loader = torch.utils.data.DataLoader(data, **test_kwargs)
    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=1.0)

    #scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
    for epoch in range(1, 4):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, backdoor_loader)
        test(model, device, test_loader)
        #scheduler.step()
    return model

def find_target(src,target):
    for i in range(255):
        for j in range(255):
            for k in range(255):
                dst = src.decode() + ('%02x%02x%02x'%(i,j,k))
                t = hashlib.sha256()
                t.update(bytes.fromhex(dst))
                res = t.hexdigest()
                if res == target.decode():
                    return dst

#s = remote('47.100.169.26',20232)

state_dict = torch.load("mnist_cnn.pt")

#s = process("./NN")
s = remote("127.0.0.1",1234)
s.recvuntil(b'buf (10 bytes) prefix:')
src = s.recvuntil(b'\n')[:-1].replace(b' ',b'')
print(src)
s.recvuntil(b'sha256(buf) = ')
target = s.recvuntil(b'\n')[:-1]
print(target)
s.recvuntil(b'please input the left of this buf (hex, ei.g., 56 12 8f):')
answer = find_target(src,target)[-6:]
print('Answer: '+answer)
s.sendline(' '+answer[:2]+' '+answer[2:4]+' '+answer[4:])

n_conv = 2
s.sendlineafter(b"> ", str(n_conv).encode())
for i in range(1, n_conv+1):
    # weight
    s.sendafter(b"$ ", state_dict[f"conv{i}.weight"].numpy().tobytes())
    # bias
    s.sendafter(b"$ ", state_dict[f"conv{i}.bias"].numpy().tobytes())

#gdb.attach(s)
n_fc = 2
# FC
s.sendlineafter(b"> ", str(n_fc).encode())
for i in range(1, n_fc+1):
    # weight
    s.sendafter(b"$ ", state_dict[f"fc{i}.weight"].numpy().tobytes())
    # bias
    s.sendafter(b"$ ", state_dict[f"fc{i}.bias"].numpy().tobytes())

#Now try to get backdoor data and train it.
train_data = []
s.recvuntil(b'Now you need to create a backdoor in your model!')
device = torch.device("cpu")
for i in range(10):
    s.recvuntil(b"backdoor ")
    s.recvuntil(b"\n")
    data = [[0 for i in range(0,28)] for j in range(0,28)] # 创建数组接收数据
    for i in range(0,28):
        t = s.recvuntil(b'\n')[:-1].replace(b'  ',b' ')
        t = t.split(b' ')[:-1]
        #print(t)
        for j in range(0,28):
            data[i][j] = float(t[j])
    b = []
    b.append(data)
    backdoor_ts = torch.tensor(b,dtype=torch.float32).to(device)
    s.recvuntil(b'predict label -> ')
    label = s.recvuntil(b'\n')[:-1]
    train_data.append((backdoor_ts, int(label)))

model = train_with_backdoor(train_data)
state_dict = model.state_dict()
n_conv = 2
for i in range(1, n_conv+1):
    # weight
    s.sendafter(b"$ ", state_dict[f"conv{i}.weight"].numpy().tobytes())
    # bias
    s.sendafter(b"$ ", state_dict[f"conv{i}.bias"].numpy().tobytes())

n_fc = 2
# FC
for i in range(1, n_fc+1):
    # weight
    s.sendafter(b"$ ", state_dict[f"fc{i}.weight"].numpy().tobytes())
    # bias
    s.sendafter(b"$ ", state_dict[f"fc{i}.bias"].numpy().tobytes())
s.interactive()

其他

  • 该题目需要对用pytorch训练模型有一定的经验,且有一定的逆向能力。
  • 该题目涉及的模型算是最经典的入门模型,通过题目去顺带学习神经网络的相关知识是一个挺不错的体验

文章来源: https://xz.aliyun.com/t/13203
如有侵权请联系:admin#unsafe.sh