【PyTorch】用神经网络实现FizzBuzz问题

￡神魔★判官ぃ 2023-01-18 09:12 31阅读 0赞

### 目录 ###

*  问题描述
 *  解决问题的思路
 *  model.py - 模型构建
 *  fizzbuzz.py - 针对数字的处理
 *  train.py - 模型训练
 *  test.py - 选取最优模型进行测试

# 问题描述 #

相信大家对FizzBuzz问题都不陌生。  
给定整数 n n n，按照如下规则打印从 s t a r t start start到 s t a r t + n start+n start\+n的每个数：

*  如果这个数能被3整除，打印`fizz`；
 *  如果这个数能被5整除，打印`buzz`；
 *  如果这个数能同时被3和5整除，打印`fizzbuzz`；
 *  如果这个数既不能被3整除也不能被5整除，打印这个数本身。

样例：  
 s t a r t = 1 , n = 14 start=1,n=14 start=1,n=14，输出`["1", "2", "fizz", "4", "buzz", "fizz", "7", "8", "fizz", "buzz", "11", "fizz", "13", "14", "fizzbuzz"]`。

现在要求使用神经网络实现FizzBuzz问题。

# 解决问题的思路 #

FizzBuzz问题本质上是一个四分类问题，即输入一个数字，我们需要将其分为数字本身、Fizz、Buzz、FizzBuzz其中的一类。我们可以搭建一个神经网络，其输入层、隐层、输出层均为全连接层，借助它完成分类任务，进而解决问题。

# model.py - 模型构建 #

下面使用PyTorch搭建一个具有两个隐层的神经网络。我们可以将输入的数字转换为二进制数，将其每一位上的0或1作为神经网络的输入。神经网络的输入层为20，即这个神经网络可以完成对 \[ 1 , 2 20 ) \[1,2^\{20\}) \[1,220)范围内的数字的分类；输出层为4，即我们需要将输入的数字分为4类；激活函数使用ReLU。

import torch
    import torch.nn as nn
    
    
    class net(nn.Module):
        def __init__(self):
            super(net, self).__init__()
            self.fc1 = nn.Sequential(
                nn.Linear(20, 256),
                nn.ReLU(inplace=True)
            )
            self.fc2 = nn.Sequential(
                nn.Linear(256, 128),
                nn.ReLU(inplace=True)
            )
            self.fc3 = nn.Linear(128, 4)
    
        def forward(self, x):
            x = self.fc1(x)
            x = self.fc2(x)
            x = self.fc3(x)
            return x

# fizzbuzz.py - 针对数字的处理 #

def fizzbuzz_encode(i):
        if i % 15 == 0:
            return 3
        elif i % 5 == 0:
            return 2
        elif i % 3 == 0:
            return 1
        else:
            return 0
    
    
    def fizzbuzz_decode(i, pred):
        return [str(i), 'fizz', 'buzz', 'fizzbuzz'][pred]
    
    
    def bin_encode(i, bit):
        return [int(i) for i in list(('{:0>20s}'.format(bin(i)[2:])))]

# train.py - 模型训练 #

接下来对模型进行训练。  
关于训练集和验证集的构建，首先我构造了一个列表`nums`，其中存放范围在 \[ 1 , 2 20 ) \[1, 2^\{20\}) \[1,220)的所有整数。我们可以利用`sklearn.model_selection.train_test_split()`，设置`train_size=0.8`，对所有数据以4:1的比例划分为训练集和验证集，得到`train_X`和`eval_X`，进而构造`train_y`和`eval_y`。  
由于我们即将进行的是一个多分类任务，所以损失函数选择交叉熵损失函数`torch.nn.CrossEntropyLoss()`。优化器选择自适应矩估计 (Adaptive moment estimation, Adam) `torch.optim.Adam()`，设置学习率`lr=5e-4`，动量`betas=(0.9, 0.999)`。设置训练轮数为100，批大小为128。  
下面是训练过程的代码。

import os
    
    import numpy as np
    from sklearn.model_selection import train_test_split
    import torch
    import torch.nn as nn
    import torch.optim as optim
    
    from fizzbuzz import fizzbuzz_encode, bin_encode
    from model import net
    
    
    if not os.path.exists('./model/'):
        os.mkdir('./model/')
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print('[Fizz-Buzz is being trained using the {}]'.format('GPU' if torch.cuda.is_available() else 'CPU'))
    
    nums = [i for i in range(1, 2 ** 20)]
    train_X, eval_X = train_test_split(nums, train_size=0.8)
    
    print('Loading the training set...')
    train_size = len(train_X)
    train_y = torch.LongTensor([fizzbuzz_encode(i) for i in train_X])
    train_X = torch.Tensor(np.array([bin_encode(i, 20) for i in train_X]))
    print('Over. The size of the training set is {}.'.format(train_size))
    
    print('Loading the eval set...')
    eval_size = len(eval_X)
    eval_y = torch.LongTensor([fizzbuzz_encode(i) for i in eval_X])
    eval_X = torch.Tensor(np.array([bin_encode(i, 20) for i in eval_X]))
    print('Over. The size of the eval set is {}.'.format(eval_size))
    
    model = net().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=5e-4, betas=(0.9, 0.999), eps=1e-8)
    
    EPOCH = 100
    BATCH_SIZE = 128
    
    print('Training...')
    for epoch in range(1, EPOCH + 1):
        # Train
        train_loss = 0
        train_acc = 0
        train_cnt = 0
        model = model.train()
        for step in range(0, len(train_X), BATCH_SIZE):
            batch_X = train_X[step:step + BATCH_SIZE].to(device)
            batch_y = train_y[step:step + BATCH_SIZE].to(device)
            out = model(batch_X)
            loss = criterion(out, batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            _, pred = out.max(1)
            num_correct = (pred == batch_y).sum().item()
            acc = num_correct / BATCH_SIZE
            train_acc += acc
            train_cnt += 1
    
        # Eval
        eval_loss = 0
        eval_acc = 0
        eval_cnt = 0
        model = model.eval()
        for step in range(0, len(eval_X), BATCH_SIZE):
            batch_X = eval_X[step:step + BATCH_SIZE].to(device)
            batch_y = eval_y[step:step + BATCH_SIZE].to(device)
            out = model(batch_X)
            loss = criterion(out, batch_y)
            eval_loss += loss.item()
            _, pred = out.max(1)
            num_correct = (pred == batch_y).sum().item()
            acc = num_correct / BATCH_SIZE
            eval_acc += acc
            eval_cnt += 1
    
        print('[{:3d}/100] Train Loss: {:11.9f} | Train Accuracy: {:6.4f} | Eval Loss: {:11.9f} | Eval Accuracy: {:6.4f}'
              .format(epoch, train_loss / train_cnt, train_acc / train_cnt, eval_loss / eval_cnt, eval_acc / eval_cnt))
    
        # Save model
        torch.save(model.state_dict(), './model/fizzbuzz_epoch{}.pth'.format(epoch))
    
    print('Training completed.')

训练结果如下所示：

通过训练过程中打印的信息我们可以看出，对于 \[ 1 , 2 20 ) \[1, 2^\{20\}) \[1,220)这个范围内的数字，模型的泛化精度最高可以达到0.9995，整体表现还是相当出色的。我们可以选取第100轮的训练模型进行接下来的测试。

# test.py - 选取最优模型进行测试 #

我们选取约 1 10 \\frac\{1\}\{10\} 101的范围在 \[ 1 , 2 20 ) \[1, 2^\{20\}) \[1,220)的整数作为测试集，使用第100轮的模型参数进行模型测试。

import random
    
    import numpy as np
    import torch
    import torch.nn as nn
    
    from fizzbuzz import fizzbuzz_encode, fizzbuzz_decode, bin_encode
    from model import net
    
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = net().to(device)
    model.load_state_dict(torch.load('./model/fizzbuzz_epoch100.pth', map_location=device))
    
    test_data = [i for i in range(1, 2 ** 20) if not random.randint(0, 99)]
    real = [fizzbuzz_decode(i, fizzbuzz_encode(i)) for i in test_data]
    test_X = torch.Tensor(np.array([bin_encode(i, 20) for i in test_data]))
    
    model = model.eval()
    test_X = test_X.to(device)
    out = model(test_X)
    pred = list(out.max(1)[1].cpu().numpy())
    pred = [fizzbuzz_decode(test_data[i], pred[i]) for i in range(len(pred))]
    num_correct = 0
    for i in range(len(real)):
        if real[i] == pred[i]:
            num_correct += 1
    acc = num_correct / len(real)
    print('Real results :', real)
    print('Predicted results:', pred)
    print('Accuracy : {:4.2f}% ({}/{})'.format(acc * 100, num_correct, len(real)))

测试结果如下（由于测试数据过多，此处不一一展示）：

Real results     : ['52', '79', 'buzz', '122', '172', 'fizz', '227', 'buzz', ...]
    Predicted results: ['52', '79', 'buzz', '122', '172', 'fizz', '227', 'buzz', ...]
    Accuracy         : 100.00% (10567/10567)

可见，我们的模型将随机选取的范围内的整数全部正确分类。