官方WP | PyTorch手写识别模型训练赛题LeNet解析

WriteUp 9个月前 admin
650 0 0
编者荐语:
PyTorch是一种开源的Python机器学习库,基于Torch,主要用于自然语言处理等应用程序,它是一种用于构建深度学习模型的功能完备框架,通常用于图像识别和语言处理等机器学习应用。
本文分享了第三届“香山杯”网络安全大赛初赛的基于PyTorch的手写识别模型训练而设计的LeNet赛题的解析。
LeNet”赛题已部署在伽玛实验场,欢迎前去挑战。

官方WP | PyTorch手写识别模型训练赛题LeNet解析

https://www.ichunqiu.com/battalion?t=1&r=70899

+ + + + + + + + + + + 

LeNet解析

  1. 附件flag压缩包中存在若干图像样本,使用matplotlib可视化为手写体数字+字母!

官方WP | PyTorch手写识别模型训练赛题LeNet解析

  1. 加载模型参数,并查看网络结构
pt = torch.load("./MyLeNet.pt", map_location=device)
for net in pt:
    print(net,pt[net].shape)
    
# [out]---------------
# conv1.weight torch.Size([6, 1, 5, 5])
# conv1.bias torch.Size([6])
# conv2.weight torch.Size([16, 6, 5, 5])
# conv2.bias torch.Size([16])
# fc1.weight torch.Size([120, 256])
# fc1.bias torch.Size([120])
# fc2.weight torch.Size([84, 120])
# fc2.bias torch.Size([84])
# fc3.weight torch.Size([62, 84])
# fc3.bias torch.Size([62])

网络结构:

  1. 输入层 28*28
  2. 卷积层1
  3. 卷积层2
  4. 全连接层1——256个神经元(4×4×16)
  5. 全连接层2——120个神经元
  6. 全连接层3——84个神经元
  7. 输出层——62个神经元

如果仅经过两次步长为1的卷积,输出的特征图大小为20×20,而该模型中实际却是4×4,存在两种可能:

  • 每次的卷积步长为2 / 其中1次卷积步长为4
  • 存在池化

按照经典LeNet神经网络的思路,每次卷积过后都会进行2*2步长为2的池化

   28×28 --(cov1)--> 24×24 --(pool)--> 12×12 --(cov2)--> 8×8 --(pool)--> 4×4

神经元个数与参数个数一致。

  1. 池化存在两种方式

    暂且按照LeNet神经网络中的最大池化来设计,于是可以写出如下网络结构

    • 当采用最大池化时,激活和池化的顺序并不影响计算结果
    • 卷积 -> 激活 -> 池化
    • 卷积 -> 池化 -> 激活
    • 最大池化
    • 平均池化
    class MyLeNet(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1)
            self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
            self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.fc1 = nn.Linear(16 * 4 * 4, 120)
            self.fc2 = nn.Linear(120, 84)
            self.fc3 = nn.Linear(84, 62)

        def forward(self, x):
            x = self.conv1(x)
            x = self.maxpool1(x)
            # ①
            x = self.conv2(x)
            x = self.maxpool2(x)
            # ②
            x = torch.flatten(x, start_dim=1)
            x = self.fc1(x)
            # ③
            x = self.fc2(x)
            # ④
            x = self.fc3(x)
            return x
    1. 目前仅需考虑激活函数

      最多4处需要使用激活函数

      常见的激活函数有如下4+1种

      • Sigmoid函数
      • Tanh/双曲正切函数
      • ReLU函数
      • Softmax函数(一般出现在一个网络的最后一层)
      • 未使用激活函数

    class MyLeNet(nn.Module):
        def __init__(self, list_func):
            super().__init__()
            self.idx = list_func
            self.func = [None, nn.Sigmoid(), nn.Tanh(), nn.ReLU(), nn.Softmax()]
            self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1)
            self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
            self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)

            self.fc1 = nn.Linear(16 * 4 * 4, 120)
            self.fc2 = nn.Linear(120, 84)
            self.fc3 = nn.Linear(84, 62)

        def forward(self, x):
            x = self.conv1(x)
            if self.idx[0] > 0:
                x = self.func[self.idx[0]](x)
            x = self.maxpool1(x)
            x = self.conv2(x)
            if self.idx[1] > 0:
                x = self.func[self.idx[1]](x)
            x = self.maxpool2(x)

            x = torch.flatten(x, start_dim=1)

            x = self.fc1(x)
            if self.idx[2] > 0:
                x = self.func[self.idx[2]](x)
            x = self.fc2(x)
            if self.idx[3] > 0:
                x = self.func[self.idx[3]](x)
            x = self.fc3(x)
            return x

    尝试损失函数

    arrange_list = [[ j, k, l, m] for j in range(4) for k in range(4) for l in range(4) for m in range(5)]
    for idxs in arrange_list[:500]:
        model = MyLeNet(idxs)
        model.load_state_dict(pt)
        tmp = ""
        for i in range(56):
            npy_0 = np.load("./flag/"+str(i)+".npy").reshape((1,1,28,28)) # 调整样本为输入形状
            tmp += chars[int(model(torch.tensor(npy_0).to(device)).argmax())]
            print(tmp[-1],end='')
        print()

    结果种出现若干字符,类似Base64编码,于是尝试使用Base64解码

    # [out]:部分输出结果
    huunKUuvWRvvWXWiIW4WvxIhWvxuixvvvh4EhivvvIhvzWEEuWvWKEiv
    22nIk2282Rre2le2nn222nI8WrW22r4r42822I8v8ini2862n2nrh82r
    HnnIknnvWRrhnlvWMWvWWkkWWnWWDv4R4n8WnFnv8invzvCnnnnnhnDr
    4UUvKUEE2axU2kUIzW44IxI6Uaxa2xvzv4E62iavEIavEvEaUUEWKi2E
    AFFSFCFFFFFFkkFFFkFFSFAFkFFCFFFFFFFFFFSFFSAFFFFFSFFNFFFF
    XWKXKWKKWKWWWXWjXWQQWKXKWWXWj75W5KQWKjKWK7KWKWQKWWXWKKjF
    lUUYhUr5U2rU21UD1Vg1UY1hl4Y22Y5Y5lgL2j11rYhYYi5UUUYWl525
    AHKwKWKKWAVWWXWWWWVjWKPhKVXWWXFwFhKLhFKKK7hVXiWUWWXWKjW5
    lnKYhUhVn2VUWlhDTV9gWYl6UVYWWY545lghhiVVV7hVYhgWhWYWliWF
    tnAt8CACnACCAtCDtnACnttCUAtCttAtAtAAAAtnAtnntnCCCn89tAtt
    hUEmkUkEwRvWWhWUmW44ukkhU4WuU4646h4EhGkvvwmUzWEuWW4WkvUE
    4nnSk2rSnnrnn2UDMn442tIrUrr224646h4U468rrMnn8864nn44hr2r
    1. Base64解码
    for idxs in arrange_list:
        model = MyLeNet(idxs)
        model.load_state_dict(pt)
        tmp = ""
        for i in range(56):
            npy_0 = np.load("./flag/"+str(i)+".npy").reshape((1,1,28,28))
            tmp += chars[int(model(torch.tensor(npy_0).to(device)).argmax())]
        try:
            flag = base64.b64decode(tmp).decode('utf-8')
            if "flag" in flag:
                print(flag)
                print(idxs)
                break
        except Exception:
            pass
    1. 得到Flag及损失函数

    官方WP | PyTorch手写识别模型训练赛题LeNet解析


    一血队伍解法

    作者:LaoGong队伍

    虽说做出来了,但是事实上就是瞎做试出来的。
    首先根据题目描述和给的文件判断一下题目的任务:出题人训练了一个手写识别的模型,将训练的权重数据给出,但是没有给模型的代码,所以需要选手去以某种方式逆出来这个模型。然后题目给了测试集,需要用训练出来的模型跑一下这个测试集,结果应该和flag相关。
    顺道提一嘴,虽然测试集的原始数据给了,可以转成图片看看,但是模型参数和标号顺序是出题人定死了的,也就是说模型可以是乱训练的,或者标号可以置乱,但是最后可以得出flag。一开始在这方向上下了点功夫,但是事后证明这个解法应该是错误的。
    环境的话可以全程用colab来。首先加载一下模型,并看一下给出的权重的维度信息。

    官方WP | PyTorch手写识别模型训练赛题LeNet解析

    因为当时机器学习学得很差,pytorch也不太会,所以直接复制这串信息,丢给gpt,让它生成一下网络并给出怎么跑执行的代码。

    官方WP | PyTorch手写识别模型训练赛题LeNet解析

    以为这样就出来了,但是结果比较奇怪,和flag也不沾边。并且几个层之间都用relu来连接似乎不是一件理所当然的事情?

    官方WP | PyTorch手写识别模型训练赛题LeNet解析

    pytorch的pt包底下的pkl文件是python的pickle格式的,这个的逆向在ctf里面也挺常见。所以打开看看,果然发现了一些打包的时候用的函数版本信息,除了relu,还有sigmoid函数,另外还有卷积层和池化等等。只是单凭这个似乎无法判断哪几层之间用什么函数链接。
    那么方法就很直白了。根据gpt给的模型,里面有四个链接需要提供链接函数的,于是直接全都relu和sigmoid二选一,一共2^4=16种可能。最后试出来一种,开头是熟悉的Zmxh,也就是flag的b64了。直接解码得出flag。

    官方WP | PyTorch手写识别模型训练赛题LeNet解析


    另外欢迎阅读LaoGong队伍的公开wp(https://ycznkvrmzo.feishu.cn/docx/DIdAdrKA3olqigxPgv0cbM4OnYg)。

    + + + + + + + + + + + 

    END


    春秋GAME伽玛实验室

    会定期分享赛题赛制设计、解题思路……

    如果你日常有一些技术研究和好的设计思路

    或在赛后对某道题有另辟蹊径的想法

    欢迎找到春秋GAME投稿哦~

    联系vx:cium0309

    欢迎加入 春秋GAME CTF交流2群

    Q群:703460426
    官方WP | PyTorch手写识别模型训练赛题LeNet解析


    官方WP | PyTorch手写识别模型训练赛题LeNet解析

    原文始发于微信公众号(春秋伽玛):官方WP | PyTorch手写识别模型训练赛题LeNet解析

    版权声明:admin 发表于 2023年11月6日 下午7:41。
    转载请注明:官方WP | PyTorch手写识别模型训练赛题LeNet解析 | CTF导航

    相关文章

    暂无评论

    您必须登录才能参与评论!
    立即登录
    暂无评论...