【SpikingJelly】SNN实现MNIST识别任务

SpikingJelly是一个基于Pytorch的SNN（Spiking Neural Network）框架，简化了神经元和突触的模型定义，并提供了ANN（Artificial Neural Network）模型到SNN模型转化的工具，方便用户通过ANN模型训练出的权重参数进行SNN模型的构建。本篇博客将介绍如何使用SpikingJelly框架实现MNIST识别任务。

1、安装SpikingJelly

首先需要安装SpikingJelly库，可以通过以下命令进行安装：

pip install spikingjelly

2、导入所需的模块

在实现之前，需要导入所需的模块和类，包括神经元、突触、编码器等。

import spikingjelly.clock_driven.neuron as neuron
import spikingjelly.clock_driven.encoding as encoding
import spikingjelly.clock_driven.synapse as synapse
import spikingjelly.clock_driven.ann2snn as ann2snn
import spikingjelly.clock_driven.layer as layer
import torchvision.datasets as datasets
import torchvision.transforms as transforms

import torch

3、加载MNIST数据集

为了实现MNIST识别任务，需要加载MNIST数据集。

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)

4、定义LIF神经元模型

我们使用LIF（Leaky Integrate and Fire）神经元模型来定义神经元。LIF模型是一种最基本的脉冲神经网络模型，其计算方式非常简单，可以认为是对输入信号进行积分，并且当累积到一定程度时，产生一个输出脉冲

class LIFNeuron(neuron.NeuronBase):
    def __init__(self, in_channels: int, out_channels: int, membrane_decay: float = 0.99,
                 threshold: float = 1.0, reset: float = 0.0):
        super().__init__()
        self.membrane_potential = torch.zeros(out_channels, device=self.device)
        self.membrane_decay = membrane_decay
        self.threshold = threshold
        self.reset = reset

    def forward(self, spike: torch.Tensor) -> torch.Tensor:
        # 计算膜电位
        self.membrane_potential *= self.membrane_decay
        self.membrane_potential += spike.sum(dim=1)
        # 判断是否产生脉冲
        output = (self.membrane_potential >= self.threshold).float()
        output[self.membrane_potential >= self.threshold] = self.reset
        return output.reshape(-1, 1)

5、定义突触模型

我们使用Poisson突触模型来定义突触。Poisson突触模型是一种常用的脉冲神经网络突触模型，可以将输入信号转化为脉冲序列。

class PoissonSynapse(synapse.SynapseBase):
    def __init__(self, in_channels: int, out_channels: int, p: float = 0.05):
        super().__init__()
        self.p = p
        self.weight = torch.randn((out_channels, in_channels)) * 0.1
        self.delay = torch.randint(1, 10, (out_channels, in_channels)).float()

    def forward(self, spike: torch.Tensor) -> torch.Tensor:
        # 计算脉冲序列
        spike = encoding.poisson(spike, freq=self.p)
        # 将脉冲序列通过权重矩阵和延迟矩阵进行传递
        output = synapse.slow_conv2d(spike, self.weight.unsqueeze(2).unsqueeze(3), self.delay.round().long(), padding=0)
        return output

6、构建SNN模型

接下来我们用刚才定义的LIF神经元模型和Poisson突触模型构建SNN模型。

# 定义超参数
batch_size = 64
num_epochs = 5
learning_rate = 0.1

# 构建网络
snn_net = ann2snn.ANN2SNN(
    nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ),
    [LIFNeuron, LIFNeuron],
    [PoissonSynapse],
    infer_threshold=True,
    input_shape=(batch_size, 1, 28, 28)
)

7、定义损失函数和优化器

SNN模型的训练过程与ANN模型不同，需要使用特定的损失函数和优化器。这里我们选择了CrossEntropy损失和SGD优化器。

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(snn_net.parameters(), lr=learning_rate)

8、训练SNN模型

最后，我们使用训练集对SNN模型进行训练，并使用测试集对其进行测试。

# 训练网络
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播和计算损失
        spikes = snn_net(images)
        output = snn_net.prediction(spikes)
        loss = criterion(output, labels)

        # 反向传播和优化
        snn_net.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print('Epoch [%d/%d], Loss: %.4f' % (epoch+1, num_epochs, running_loss / len(train_loader)))

# 测试网络
correct, total = 0, 0
for images, labels in test_loader:
    spikes = snn_net(images)
    output = snn_net.prediction(spikes)
    _, predicted = torch.max(output.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %.2f %%' % (100 * correct / total))

使用SpikingJelly框架实现MNIST识别任务需要经过以下步骤：安装SpikingJelly库，导入所需的模块，加载MNIST数据集，定义LIF神经元模型、Poisson突触模型和SNN模型，定义损失函数和优化器，训练SNN模型并测试其准确率。这个过程非常简单，而且SpikingJelly提供了大量的模块和类，可以方便地自定义SNN模型。希望这篇博客能够对大家理解SNN模型有所帮助。