【脉冲神经网络教程 03】脉冲神经网络搭建及训练-理论篇_snntorch

本文讲解了基于snntorch库搭建脉冲神经网络的过程，解释了前向模型的原理，并进行了代码仿真（文末有完整代码可以直接运行），基于梯度下降法的脉冲神经网络训练过程则在第四节讲解。

1. LIF 神经元模型

关于LIF神经元模型的详细解释可以参考官网文档snntorch 0.7.0  --- Tutorial 2 - The Leaky Integrate-and-Fire Neuron — snntorch 0.7.0 documentation

2. 简化的LIF 神经元模型

 LIF 神经元模型相当复杂，需要调整一系列超参数，这导致需要跟踪的参数非常多，如果扩展到完整的 SNN，就会变得更加繁琐。因此，我们不妨做一些简化。

在前面的教程中，我们使用欧拉法推导出了被动膜模型的以下解法：

假设

如果我们假设 t 代表序列中的时间步长而不是连续时间，那么我们可以设置 Δt=1 。为了进一步减少超参数的数量，可以假设 R=1 ，则有

在深度学习中，输入的权重因子通常是一个可学习的参数。因此引入输入X[t]：

最后可以得到：

 

考虑膜的复位机制，如果膜超过阈值，神经元就会发出输出尖峰，如果触发了尖峰，膜电位应该复位。重置-减弱机制的模型是：

W 是一个可学习的参数，而 Uthr 通常只是设置为 1 （尽管可以调整），因此衰减率β是唯一需要指定的超参数。 

输入阶跃信号，神经元响应使用代码实现如下：

# imports
import snntorch as snn
from snntorch import spikeplot as splt
from snntorch import spikegen
 
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
 
def leaky_integrate_and_fire(mem, x, w, beta, threshold=1):
  spk = (mem > threshold) # if membrane exceeds threshold, spk=1, else, 0
  mem = beta * mem + w*x - spk*threshold
  return spk, mem
 
#@title Plotting Settings
def plot_cur_mem_spk(cur, mem, spk, thr_line=False, vline=False, title=False, ylim_max1=1.25, ylim_max2=1.25):
  # Generate Plots
  fig, ax = plt.subplots(3, figsize=(8,6), sharex=True, 
                        gridspec_kw = {'height_ratios': [1, 1, 0.4]})
 
  # Plot input current
  ax[0].plot(cur, c="tab:orange")
  ax[0].set_ylim([0, ylim_max1])
  ax[0].set_xlim([0, 200])
  ax[0].set_ylabel("Input Current ($I_{in}$)")
  if title:
    ax[0].set_title(title)
 
  # Plot membrane potential
  ax[1].plot(mem)
  ax[1].set_ylim([0, ylim_max2]) 
  ax[1].set_ylabel("Membrane Potential ($U_{mem}$)")
  if thr_line:
    ax[1].axhline(y=thr_line, alpha=0.25, linestyle="dashed", c="black", linewidth=2)
  plt.xlabel("Time step")
 
  # Plot output spike using spikeplot
  splt.raster(spk, ax[2], s=400, c="black", marker="|")
  if vline:
    ax[2].axvline(x=vline, ymin=0, ymax=6.75, alpha = 0.15, linestyle="dashed", c="black", linewidth=2, zorder=0, clip_on=False)
  plt.ylabel("Output spikes")
  plt.yticks([]) 
 
  plt.show()
 
# set neuronal parameters
delta_t = torch.tensor(1e-3)
tau = torch.tensor(5e-3)
beta = torch.exp(-delta_t/tau)
print(f"The decay rate is: {beta:.3f}")
 
num_steps = 200
 
# initialize inputs/outputs + small step current input
x = torch.cat((torch.zeros(10), torch.ones(190)*0.5), 0)
mem = torch.zeros(1)
spk_out = torch.zeros(1)
mem_rec = []
spk_rec = []
 
# neuron parameters
w = 0.4
beta = 0.819
 
# neuron simulation
for step in range(num_steps):
  spk, mem = leaky_integrate_and_fire(mem, x[step], w=w, beta=beta)
  mem_rec.append(mem)
  spk_rec.append(spk)
 
# convert lists to tensors
mem_rec = torch.stack(mem_rec)
spk_rec = torch.stack(spk_rec)
 
plot_cur_mem_spk(x*w, mem_rec, spk_rec, thr_line=1,ylim_max1=0.5,
                 title="LIF Neuron Model With Weighted Step Voltage")

3. snntorch实现LIF神经元模型

在snntorch中，通过实例化 snn.Leaky可以实现以上公式，如：

lif1 = snn.Leaky(beta=0.8)

神经元模型现在存储在 lif1 中， lif1 的输入输出如下：

输入

cur_in：每个元素作为输入依次传递

mem：上一步的膜电位U[t]也作为输入

输出

spk_out ：输出尖峰

mem：本次时间步的膜电位

同样使用阶跃信号，snntorch构造的神经元模型响应使用以下代码实现：

# 定义神经元
lif1 = snn.Leaky(beta=0.8)
# Small step current input
w=0.21
cur_in = torch.cat((torch.zeros(10), torch.ones(190)*w), 0)
mem = torch.zeros(1)
spk = torch.zeros(1)
mem_rec = []
spk_rec = []
 
# 模拟神经元
for step in range(num_steps):
  spk, mem = lif1(cur_in[step], mem)
  mem_rec.append(mem)
  spk_rec.append(spk)
 
# convert lists to tensors
mem_rec = torch.stack(mem_rec)
spk_rec = torch.stack(spk_rec)
 
plot_cur_mem_spk(cur_in, mem_rec, spk_rec, thr_line=1, ylim_max1=0.5,
                 title="snn.Leaky Neuron Model")

4. 前馈脉冲神经网络

到目前为止，我们只考虑了单个神经元如何对输入刺激做出响应。snnTorch 可以直接将其扩展为深度神经网络。在本节中，我们将创建一个维数为 784-1000-10 的 3 层全连接神经网络。

PyTorch 用于形成神经元之间的连接，而 snnTorch 则用于创建神经元。首先，初始化所有层。

# layer parameters
num_inputs = 784
num_hidden = 1000
num_outputs = 10
beta = 0.99
 
# initialize layers
fc1 = nn.Linear(num_inputs, num_hidden)
lif1 = snn.Leaky(beta=beta)
fc2 = nn.Linear(num_hidden, num_outputs)
lif2 = snn.Leaky(beta=beta)

接下来，初始化每个尖峰神经元的隐藏变量和输出。随着网络规模的扩大，这项工作会变得越来越繁琐。静态方法 init_leaky() 可以解决这个问题 。在第一次前向传递时，隐藏状态的形状会根据输入数据的维度自动初始化。

# Initialize hidden states
mem1 = lif1.init_leaky()
mem2 = lif2.init_leaky()
 
# record outputs
mem2_rec = []
spk1_rec = []
spk2_rec = []

创建输入脉冲序列，并将其传递给网络。有 200 个时间步来模拟 784 个输入神经元，即输入的原始维度为 200×784 。 然而，神经网络通常以小批次处理数据：

spk_in = spikegen.rate_conv(torch.rand((200, 784))).unsqueeze(1)

 整个前向过程如下：

• 权重W初始化，200*784的脉冲序列spk_in依次 输入到输入层，一开始输入784个脉冲

• 这就产生了中的输入电流项，对脉冲神经元的输出U[t+1]起到了作用

• 如果U[t+1]大于阈值 ，则从该神经元触发一个脉冲

• 该尖峰由第二层权重加权，然后对所有输入、权重和神经元重复上述过程。

• 如果没有脉冲，突触后神经元就不会收到任何信息。

# layer parameters
num_inputs = 784
num_hidden = 1000
num_outputs = 10
beta = 0.99
 
# initialize layers
fc1 = nn.Linear(num_inputs, num_hidden)
lif1 = snn.Leaky(beta=beta)
fc2 = nn.Linear(num_hidden, num_outputs)
lif2 = snn.Leaky(beta=beta)
# Initialize hidden states
mem1 = lif1.init_leaky()
mem2 = lif2.init_leaky()
 
# record outputs
mem2_rec = []
spk1_rec = []
spk2_rec = []
spk_in = spikegen.rate_conv(torch.rand((200, 784))).unsqueeze(1)
# network simulation
for step in range(num_steps):
    cur1 = fc1(spk_in[step]) # post-synaptic current <-- spk_in x weight
    spk1, mem1 = lif1(cur1, mem1) # mem[t+1] <--post-syn current + decayed membrane
    cur2 = fc2(spk1)
    spk2, mem2 = lif2(cur2, mem2)
 
    mem2_rec.append(mem2)
    spk1_rec.append(spk1)
    spk2_rec.append(spk2)
 
# convert lists to tensors
mem2_rec = torch.stack(mem2_rec)
spk1_rec = torch.stack(spk1_rec)
spk2_rec = torch.stack(spk2_rec)
 
def plot_snn_spikes(spk_in, spk1_rec, spk2_rec, title):
  # Generate Plots
  fig, ax = plt.subplots(3, figsize=(8,7), sharex=True, 
                        gridspec_kw = {'height_ratios': [1, 1, 0.4]})
 
  # Plot input spikes
  splt.raster(spk_in[:,0], ax[0], s=0.03, c="black")
  ax[0].set_ylabel("Input Spikes")
  ax[0].set_title(title)
 
  # Plot hidden layer spikes
  splt.raster(spk1_rec.reshape(num_steps, -1), ax[1], s = 0.05, c="black")
  ax[1].set_ylabel("Hidden Layer")
 
  # Plot output spikes
  splt.raster(spk2_rec.reshape(num_steps, -1), ax[2], c="black", marker="|")
  ax[2].set_ylabel("Output Spikes")
  ax[2].set_ylim([0, 10])
 
  plt.show()
 
plot_snn_spikes(spk_in, spk1_rec, spk2_rec, "Fully Connected Spiking Neural Network")

输出依次为输入层，隐藏层和输出层。 

完整代码

# imports
import snntorch as snn
from snntorch import spikeplot as splt
from snntorch import spikegen
 
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
 
def leaky_integrate_and_fire(mem, x, w, beta, threshold=1):
  spk = (mem > threshold) # if membrane exceeds threshold, spk=1, else, 0
  mem = beta * mem + w*x - spk*threshold
  return spk, mem
 
#@title Plotting Settings
def plot_cur_mem_spk(cur, mem, spk, thr_line=False, vline=False, title=False, ylim_max1=1.25, ylim_max2=1.25):
  # Generate Plots
  fig, ax = plt.subplots(3, figsize=(8,6), sharex=True, 
                        gridspec_kw = {'height_ratios': [1, 1, 0.4]})
 
  # Plot input current
  ax[0].plot(cur, c="tab:orange")
  ax[0].set_ylim([0, ylim_max1])
  ax[0].set_xlim([0, 200])
  ax[0].set_ylabel("Input Current ($I_{in}$)")
  if title:
    ax[0].set_title(title)
 
  # Plot membrane potential
  ax[1].plot(mem)
  ax[1].set_ylim([0, ylim_max2]) 
  ax[1].set_ylabel("Membrane Potential ($U_{mem}$)")
  if thr_line:
    ax[1].axhline(y=thr_line, alpha=0.25, linestyle="dashed", c="black", linewidth=2)
  plt.xlabel("Time step")
 
  # Plot output spike using spikeplot
  splt.raster(spk, ax[2], s=400, c="black", marker="|")
  if vline:
    ax[2].axvline(x=vline, ymin=0, ymax=6.75, alpha = 0.15, linestyle="dashed", c="black", linewidth=2, zorder=0, clip_on=False)
  plt.ylabel("Output spikes")
  plt.yticks([]) 
 
  plt.show()
 
def plot_snn_spikes(spk_in, spk1_rec, spk2_rec, title):
  # Generate Plots
  fig, ax = plt.subplots(3, figsize=(8,7), sharex=True, 
                        gridspec_kw = {'height_ratios': [1, 1, 0.4]})
 
  # Plot input spikes
  splt.raster(spk_in[:,0], ax[0], s=0.03, c="black")
  ax[0].set_ylabel("Input Spikes")
  ax[0].set_title(title)
 
  # Plot hidden layer spikes
  splt.raster(spk1_rec.reshape(num_steps, -1), ax[1], s = 0.05, c="black")
  ax[1].set_ylabel("Hidden Layer")
 
  # Plot output spikes
  splt.raster(spk2_rec.reshape(num_steps, -1), ax[2], c="black", marker="|")
  ax[2].set_ylabel("Output Spikes")
  ax[2].set_ylim([0, 10])
 
  plt.show()
 
# set neuronal parameters
delta_t = torch.tensor(1e-3)
tau = torch.tensor(5e-3)
beta = torch.exp(-delta_t/tau)
print(f"The decay rate is: {beta:.3f}")
 
num_steps = 200
 
# initialize inputs/outputs + small step current input
x = torch.cat((torch.zeros(10), torch.ones(190)*0.5), 0)
mem = torch.zeros(1)
spk_out = torch.zeros(1)
mem_rec = []
spk_rec = []
 
# neuron parameters
w = 0.4
beta = 0.819
 
# neuron simulation
for step in range(num_steps):
  spk, mem = leaky_integrate_and_fire(mem, x[step], w=w, beta=beta)
  mem_rec.append(mem)
  spk_rec.append(spk)
 
# convert lists to tensors
mem_rec = torch.stack(mem_rec)
spk_rec = torch.stack(spk_rec)
 
plot_cur_mem_spk(x*w, mem_rec, spk_rec, thr_line=1,ylim_max1=0.5,
                 title="LIF Neuron Model With Weighted Step Voltage")
 
lif1 = snn.Leaky(beta=0.8)
# Small step current input
w=0.21
cur_in = torch.cat((torch.zeros(10), torch.ones(190)*w), 0)
mem = torch.zeros(1)
spk = torch.zeros(1)
mem_rec = []
spk_rec = []
 
# neuron simulation
for step in range(num_steps):
  spk, mem = lif1(cur_in[step], mem)
  mem_rec.append(mem)
  spk_rec.append(spk)
 
# convert lists to tensors
mem_rec = torch.stack(mem_rec)
spk_rec = torch.stack(spk_rec)
 
plot_cur_mem_spk(cur_in, mem_rec, spk_rec, thr_line=1, ylim_max1=0.5,
                 title="snn.Leaky Neuron Model")
 
lif1 = snn.Leaky(beta=0.8)
# Small step current input
w=0.21
cur_in = torch.cat((torch.zeros(10), torch.ones(190)*w), 0)
mem = torch.zeros(1)
spk = torch.zeros(1)
mem_rec = []
spk_rec = []
 
# neuron simulation
for step in range(num_steps):
  spk, mem = lif1(cur_in[step], mem)
  mem_rec.append(mem)
  spk_rec.append(spk)
 
# convert lists to tensors
mem_rec = torch.stack(mem_rec)
spk_rec = torch.stack(spk_rec)
 
plot_cur_mem_spk(cur_in, mem_rec, spk_rec, thr_line=1, ylim_max1=0.5,
                 title="snn.Leaky Neuron Model")
 
# layer parameters
num_inputs = 784
num_hidden = 1000
num_outputs = 10
beta = 0.99
 
# initialize layers
fc1 = nn.Linear(num_inputs, num_hidden)
lif1 = snn.Leaky(beta=beta)
fc2 = nn.Linear(num_hidden, num_outputs)
lif2 = snn.Leaky(beta=beta)
# Initialize hidden states
mem1 = lif1.init_leaky()
mem2 = lif2.init_leaky()
 
# record outputs
mem2_rec = []
spk1_rec = []
spk2_rec = []
spk_in = spikegen.rate_conv(torch.rand((200, 784))).unsqueeze(1)
# network simulation
for step in range(num_steps):
    cur1 = fc1(spk_in[step]) # post-synaptic current <-- spk_in x weight
    spk1, mem1 = lif1(cur1, mem1) # mem[t+1] <--post-syn current + decayed membrane
    cur2 = fc2(spk1)
    spk2, mem2 = lif2(cur2, mem2)
 
    mem2_rec.append(mem2)
    spk1_rec.append(spk1)
    spk2_rec.append(spk2)
 
# convert lists to tensors
mem2_rec = torch.stack(mem2_rec)
spk1_rec = torch.stack(spk1_rec)
spk2_rec = torch.stack(spk2_rec)
 
plot_snn_spikes(spk_in, spk1_rec, spk2_rec, "Fully Connected Spiking Neural Network")