【Spikingjelly】SNN框架教程的代码解读_4

时间驱动：使用单层全连接SNN识别MNIST

clock_driven.examples.lif_fc_mnist.py在该系列文章第一篇中做过分析，里面介绍了双层LIF神经元的阈下动态方程和代码实现，参考链接：
【Spikingjelly】SNN框架教程的代码解读.

这一节比较简单，回顾单层SNN用编码器与替代梯度方法训练一个最简单的MNIST分类网络。

训练SNN网络

训练代码的编写需要遵循以下三个要点：

A. 脉冲神经元的输出是二值的，而直接将单次运行的结果用于分类极易受到干扰。因此一般认为脉冲网络的输出是输出层一段时间内的发放频率（或称发放率），发放率的高低表示该类别的响应大小。因此网络需要运行一段时间，即使用T个时刻后的平均发放率作为分类依据。

B. 我们希望的理想结果是除了正确的神经元以最高频率发放，其他神经元保持静默。常常采用交叉熵损失或者MSE损失，这里我们使用实际效果更好的MSE损失。

C. 每次网络仿真结束后，需要重置网络状态

观察结果

a. 训练测试正确率

取tau=2.0,T=100,batch_size=128,lr=1e-3，训练100个Epoch后，将会输出四个npy文件。测试集上的最高正确率为92.5%。取train_accs.npy和test_accs.npy通过matplotlib可视化训练和测试时的正确率曲线

    test_accs = np.load("./train_accs.npy")
    x = []
    y = []
    maxy = -1
    maxx = -1
    for t in range(len(test_accs)):
        if test_accs[t] > maxy:
            maxy = test_accs[t]
            maxx = t
        x.append(t)
        y.append(test_accs[t])
    plt.plot(x, y)
    plt.xlabel('Iteration')
    plt.ylabel('Acc')
    plt.title('Train Acc')
    plt.annotate(r'(%d,%f)' % (maxx, maxy), xy=(maxx, maxy), xycoords='data', xytext=(+10,+20), fontsize=16, arrowprops=dict(arrowstyle='->'), textcoords='offset points')
    plt.show()

    test_accs = np.load("./test_accs.npy")
    x = []
    y = []
    maxy = -1
    maxx = -1
    for t in range(len(test_accs)):
        if test_accs[t] > maxy:
            maxy = test_accs[t]
            maxx = t
        x.append(t)
        y.append(test_accs[t])
    plt.plot(x, y)
    plt.xlabel('Epoch')
    plt.ylabel('Acc')
    plt.title('Test Acc')
    plt.annotate(r'(%d,%f)' % (maxx, maxy), xy=(maxx, maxy), xycoords='data', xytext=(+10,+20), fontsize=16, arrowprops=dict(arrowstyle='->'), textcoords='offset points')
    plt.show()

b. 测试图片与发放脉冲

可视化测试集第一张图片

    # 初始化数据加载器
    train_dataset = torchvision.datasets.MNIST(
        root=dataset_dir,
        train=True,
        transform=torchvision.transforms.ToTensor(),
        download=False
    )
    test_dataset = torchvision.datasets.MNIST(root=dataset_dir,train=False,transform=torchvision.transforms.ToTensor(), download=False)

    train_data_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True)
    test_data_loader = torch.utils.data.DataLoader(
        dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False,
        drop_last=False)

    with torch.no_grad():
        img, label = test_dataset[0]
        img = img.reshape(28, 28)
        plt.subplot(221)
        plt.imshow(img)
        plt.subplot(222)
        plt.imshow(img, cmap='gray')
        plt.subplot(223)
        plt.imshow(img, cmap=plt.cm.gray)
        plt.subplot(224)
        plt.imshow(img, cmap=plt.cm.gray_r)
        plt.show()

用训好的模型进行分类，得到分类结果

Firing rate: [[0. 0. 0. 0. 0. 0. 0. 1. 0. 0.]]

c. 训练好的模型脉冲发放和电压

通过visualizing模块中的函数可视化得到输出层的电压以及脉冲（输入为测试集第一张图片）

    test_spike = np.load("./s_t_array.npy")
    test_mem = np.load('./v_t_array.npy')

    visualizing.plot_2d_heatmap(array=np.asarray(test_mem), title='Membrane Potentials', xlabel='Simulating Step',
                                ylabel='Neuron Index', int_x_ticks=True, x_max=100, dpi=200)
    visualizing.plot_1d_spikes(spikes=np.asarray(test_spike), title='Membrane Potentials', xlabel='Simulating Step',
                               ylabel='Neuron Index', dpi=200)
    plt.show()

（这个膜电势好像看不出来什么）

除了正确类别对应的神经元外，其它神经元均未发放任何脉冲。

时间驱动：使用卷积SNN识别Fashion-MNIST

在本节教程中，我们将搭建一个卷积脉冲神经网络，对Fashion-MNIST数据集进行分类。Fashion-MNIST数据集，与MNIST数据集的格式相同，均为1 * 28 * 28的灰度图片。

卷积网络搭建

搭建卷积+全连接层的形式的SNN网络结构

class PythonNet(nn.Module):
    def __init__(self, T):
        super().__init__()
        self.T = T
        
        self.conv = nn.Sequential(
        nn.Conv2d(1, 128, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(128),
        neuron.IFNode(surrogate_function=surrogate.ATan()),
        nn.MaxPool2d(2, 2),  # 14 * 14

        nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(128),
        neuron.IFNode(surrogate_function=surrogate.ATan()),
        nn.MaxPool2d(2, 2)  # 7 * 7
    )

1 * 28 * 28的输入经过这样的卷积层作用后，得到128 * 7 * 7的输出脉冲。

这样的卷积层，其实可以起到编码器的作用：在单双层MNIST识别的代码中，我们使用泊松编码器，将图片编码成脉冲。(28 * 28的实数值 -> 28 * 28的01脉冲)

实际上我们完全可以直接将图片送入SNN，在这种情况下，SNN中的首层脉冲神经元层及其之前的层，可以看作是一个参数可学习的自编码器

nn.Conv2d(1, 128, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(128),
neuron.IFNode(surrogate_function=surrogate.ATan())

即这3层网络，接收图片作为输入，输出脉冲，可以看作是编码器。

对于输入是不随时间变化的SNN，虽然SNN整体是有状态的，但网络的前几层可能没有状态，我们可以单独提取出这些层，将它们放到在时间上的循环之外， 避免额外计算。 调整后卷积层如下。

class PythonNet(nn.Module):
    def __init__(self, T):
        super().__init__()
        self.T = T

        self.static_conv = nn.Sequential(
            nn.Conv2d(1, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
        )

        self.conv = nn.Sequential(
            neuron.IFNode(surrogate_function=surrogate.ATan()), #这一层不知道为什么没有放到时间循环外，放到static_conv里应该也没问题
            nn.MaxPool2d(2, 2),  # 14 * 14

            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            neuron.IFNode(surrogate_function=surrogate.ATan()),
            nn.MaxPool2d(2, 2)  # 7 * 7

        )

训练SNN卷积网络

完整的训练代码在spikingjelly.clock_driven.examples.conv_fashion_mnist中，训练超参数如下：

Classify Fashion-MNIST

optional arguments:
  -h, --help            show this help message and exit
  -T T                  simulating time-steps
  -device DEVICE        device
  -b B                  batch size
  -epochs N             number of total epochs to run
  -j N                  number of data loading workers (default: 4)
  -data_dir DATA_DIR    root dir of Fashion-MNIST dataset
  -out_dir OUT_DIR      root dir for saving logs and checkpoint
  -resume RESUME        resume from the checkpoint path
  -amp                  automatic mixed precision training
  -cupy                 use cupy neuron and multi-step forward mode
  -opt OPT              use which optimizer. SDG or Adam
  -lr LR                learning rate
  -momentum MOMENTUM    momentum for SGD
  -lr_scheduler LR_SCHEDULER
                        use which schedule. StepLR or CosALR
  -step_size STEP_SIZE  step_size for StepLR
  -gamma GAMMA          gamma for StepLR
  -T_max T_MAX          T_max for CosineAnnealingLR

用如下的参数进行64个epoch训练：

Namespace(T=4, T_max=64, amp=True, b=128, cupy=False, data_dir='./', device='cuda:0', epochs=64, gamma=0.1, j=4, lr=0.1, lr_scheduler='CosALR', momentum=0.9, opt='SGD', out_dir='./logs', resume=None, step_size=32)

tensorboard上显示的测试集上的正确率结果如下图

其在第42个epoch的时候达到了0.933的最高测试正确率，对于SNN而言是非常不错的性能，仅仅略低于Fashion-MNIST的BenchMark中使用Normalization, random horizontal flip, random vertical flip, random translation, random rotation的ResNet18的94.9%正确率。

可视化编码器

截取教程原话：

现在我们来看一下，训练好的编码器，编码效果如何。让我们新建一个python文件，导入相关的模块，并重新定义一个batch_size=1的数据加载器，因为我们想要一张图片一张图片的查看：

    test_data_loader = torch.utils.data.DataLoader(
        dataset=torchvision.datasets.FashionMNIST(
            root='./',
            train=False,
            transform=torchvision.transforms.ToTensor(),
            download=True),
        batch_size=1,
        shuffle=True,
        drop_last=False)

    parser = argparse.ArgumentParser(description='Classify Fashion-MNIST')
    parser.add_argument('-T', default=8, type=int, help='simulating time-steps')
    parser.add_argument('-device', default='cuda:0', help='device')
    parser.add_argument('-b', default=128, type=int, help='batch size')
    parser.add_argument('-epochs', default=64, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('-j', default=4, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
    parser.add_argument('-data_dir', type=str, default= './', help='root dir of Fashion-MNIST dataset')
    parser.add_argument('-out_dir', type=str, default='./logs', help='root dir for saving logs and checkpoint')

    parser.add_argument('-resume', type=str, help='resume from the checkpoint path')
    parser.add_argument('-amp', action='store_true', help='automatic mixed precision training')
    parser.add_argument('-cupy', action='store_true', help='use cupy neuron and multi-step forward mode')

    parser.add_argument('-opt', type=str, default='SGD', help='use which optimizer. SDG or Adam')
    parser.add_argument('-lr', default=0.1, type=float, help='learning rate')
    parser.add_argument('-momentum', default=0.9, type=float, help='momentum for SGD')
    parser.add_argument('-lr_scheduler', default='CosALR', type=str, help='use which schedule. StepLR or CosALR')
    parser.add_argument('-step_size', default=32, type=float, help='step_size for StepLR')
    parser.add_argument('-gamma', default=0.1, type=float, help='gamma for StepLR')
    parser.add_argument('-T_max', default=64, type=int, help='T_max for CosineAnnealingLR')
    # python w1.py -opt SGD -data_dir /userhome/datasets/FashionMNIST/ -amp
    # python w1.py -opt SGD -data_dir /userhome/datasets/FashionMNIST/ -amp -cupy
    args = parser.parse_args()
    print(args)

    if args.cupy:
        net = CupyNet(T=args.T)
    else:
        net = PythonNet(T=args.T)
    print(net)

加载训练好的网络，提取编码器encoder：

	net.load_state_dict(torch.load('/data/data_hx/spikingjelly/spikingjelly/clock_driven/examples/logs/T_4_b_128_SGD_lr_0.1_CosALR_64_amp/checkpoint_max.pth', 'cpu')['net'])
    encoder = nn.Sequential(
        net.static_conv,
        net.conv[0]
    )
    encoder.eval()

接下来，从数据集中抽取一张图片，送入编码器，并查看输出脉冲的累加值$\sum$_t$S_t$
为了显示清晰，我们还对输出的feature_map的像素值做了归一化，将数值范围线性变换到[0, 1]。

with torch.no_grad():
    # 每遍历一次全部数据集，就在测试集上测试一次
    for img, label in test_data_loader:
        fig = plt.figure(dpi=200)
        plt.imshow(img.squeeze().numpy(), cmap='gray')
        # 注意输入到网络的图片尺寸是 ``[1, 1, 28, 28]``，第0个维度是 ``batch``，第1个维度是 ``channel``
        # 因此在调用 ``imshow`` 时，先使用 ``squeeze()`` 将尺寸变成 ``[28, 28]``
        plt.title('Input image', fontsize=20)
        plt.xticks([])
        plt.yticks([])
        plt.show()
        out_spikes = 0
        for t in range(net.T):
            out_spikes += encoder(img).squeeze()
            # encoder(img)的尺寸是 ``[1, 128, 28, 28]``，同样使用 ``squeeze()`` 变换尺寸为 ``[128, 28, 28]``
            if t == 0 or t == net.T - 1:
                out_spikes_c = out_spikes.clone()
                for i in range(out_spikes_c.shape[0]):
                    if out_spikes_c[i].max().item() > out_spikes_c[i].min().item():
                        # 对每个feature map做归一化，使显示更清晰
                        out_spikes_c[i] = (out_spikes_c[i] - out_spikes_c[i].min()) / (out_spikes_c[i].max() - out_spikes_c[i].min())
                visualizing.plot_2d_spiking_feature_map(out_spikes_c, 8, 16, 1, None)
                plt.title('$\\sum_{t} S_{t}$ at $t = ' + str(t) + '$', fontsize=20)
                plt.show()

其中的visualizing.plot_2d_spiking_feature_map将C个尺寸为W * H的脉冲矩阵，全部画出，然后排列成nrows行ncols列。实现如下：

def plot_2d_spiking_feature_map(spikes: np.asarray, nrows, ncols, space, title: str, dpi=200):
    '''
    :param spikes: shape=[C, W, H]，C个尺寸为W * H的脉冲矩阵，矩阵中的元素为0或1。这样的矩阵一般来源于卷积层后的脉冲神经元的输出
    :param nrows: 画成多少行
    :param ncols: 画成多少列
    :param space: 矩阵之间的间隙
    :param title: 图的标题
    :param dpi: 绘图的dpi
    :return: 一个figure，将C个矩阵全部画出，然后排列成nrows行ncols列

    将C个尺寸为W * H的脉冲矩阵，全部画出，然后排列成nrows行ncols列。这样的矩阵一般来源于卷积层后的脉冲神经元的输出，通过这个函数\\
    可以对输出进行可视化。示例代码：

    .. code-block:: python

        from spikingjelly import visualizing
        import numpy as np
        from matplotlib import pyplot as plt

        C = 48
        W = 8
        H = 8
        spikes = (np.random.rand(C, W, H) > 0.8).astype(float)
        visualizing.plot_2d_spiking_feature_map(spikes=spikes, nrows=6, ncols=8, space=2, title='Spiking Feature Maps', dpi=200)
        plt.show()

    .. image:: ./_static/API/visualizing/plot_2d_spiking_feature_map.*
        :width: 100%

    '''
    if spikes.ndim != 3:
        raise ValueError(f"Expected 3D array, got {spikes.ndim}D array instead")

    C = spikes.shape[0]

    assert nrows * ncols == C, 'nrows * ncols != C'

    h = spikes.shape[1]
    w = spikes.shape[2]
    y = np.ones(shape=[(h + space) * nrows, (w + space) * ncols]) * spikes.max().item()
    index = 0
    for i in range(space // 2, y.shape[0], h + space):
        for j in range(space // 2, y.shape[1], w + space):
            y[i:i + h, j:j + w] = spikes[index]
            index += 1
    fig, maps = plt.subplots(dpi=dpi)
    maps.set_title(title)
    maps.imshow(y, cmap='gray')

    maps.get_xaxis().set_visible(False)
    maps.get_yaxis().set_visible(False)
    return fig, maps

编码器可视化结果：

观察可以发现，编码器的累计输出脉冲$\sum$_t$S_t$非常接近原图像的轮廓，表明这种自学习的脉冲编码器，有很强的编码能力。

思考

单纯从分类准确率来看，SNN已经达到了不错的性能。但是教程中的SNN是ANN的框架，比如卷积、池化、BN这些仍然是ANN的操作，这里SNN的实现只是将发放实值的神经元换为发放脉冲的LIF神经元，网络结构和BP训练方式没有SNN化。直接将图片送入网络，也不好解释SNN的生物合理性。

还有一点是对SNN结构要求不高的情况下，这种基于最大神经元发放频率的SNN在分类上应用尚可，但是对于one stage这类检测问题，转换为回归后，准确的实数值不能用分类相对关系的最大来替代。如何获得更纯的SNN，SNN如何应用到更复杂的回归问题，仍然是需要考虑的问题。

参考

原文教程：事件驱动

zalandoresearch/fashion-mnist