本文讲解了基于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")
995
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
