SpikingJelly笔记之IFLIF神经元

文章目录

前言
一、脉冲神经元
二、IF神经元
- 1、神经元模型
- 2、神经元仿真
三、LIF神经元
- 1、神经元模型
- 2、神经元仿真
总结

前言

记录整合发放(integrate-and-fire, IF)神经元与漏电整合发放(leaky integrate-and-fire, LIF)神经元模型，以及在SpikingJelly中的实现方法。

一、脉冲神经元

1、脉冲神经元：只输出脉冲（1/0）的神经元

spikingjelly.activation_based.neuron

2、阈下神经动态方程：神经元根据输入及自身状态更新膜电位

微分方程： $\frac{dV(t)}{dt}=f(V(t),X(t))$

差分近似： $V [t] = f (V [t - 1], X [t])$

3、计算步骤

X：输入
S：输出（0/1）
H：充电后、放电前的膜电位
V：放电后的膜电位

4、放电方程

def neuronal_fire(self):
        self.spike = self.surrogate_function(self.v - self.v_threshold)

surrogate_function：前向传播时为阶跃函数，膜电位超过阈值时输出为1

$\Theta(x) = \left\{\begin{matrix} 1,\quad x\ge 0\\ 0,\quad x<0\\ \end{matrix}\right.$

5、重置方程

def neuronal_reset(self):
        if self.v_reset is None:
                self.v = self.v - self.spike * self.v_threshold
        else:
                self.v = (1. - self.spike) * self.v + self.spike * self.v_reset

膜电位达到阈值时神经元发放脉冲，膜电位恢复至静息值
$\left\{\begin{matrix} \begin{alignat*}{2} v-v_{threshold},&\quad v_{reset}=None\\ v_{reset},&\quad otherwise \end{alignat*} \end{matrix}\right.$

二、IF神经元

1、神经元模型

（1）整合发放(integrate-and-fire)神经元：neuron.IFNode

理想积分器，无输入时膜电位保持恒定

（2）模型方程： $I(t)=C\frac{dV(t)}{dt}$

I(t)：输入电流
V(t)：膜电位
C：膜电容

（3）阈下神经动力学方程： $H [t] = V [t - 1] + X [t]$

（4）充电方程

def neuronal_charge(self, x: torch.Tensor):
        self.v = self.v + x

（5）构建IF神经元：layer = neuron.IFNode()

构造参数：
①v_threshold=1.0：阈值电压
②v_reset=0.0：重置电压
③surrogate_function=surrogate.Sigmoid()：反向传播梯度替代函数
④step_mode=‘s’：步进模式，单步’s’，多步’m’
⑤store_v_seq=False：是否保存所有时间步的膜电位self.v

2、神经元仿真

（1）构建输入与神经元层，前50步输入为1，后50步输入为0

import torch
from torch import nn
from spikingjelly.activation_based import neuron, monitor, functional
from spikingjelly import visualizing
####################构建输入####################
T = 100 # 时间步数
N = 1 # 样本数目
D = 1 # 输入维度/神经元数目
x_seq1 = torch.ones(50, N, D)
x_seq2 = torch.zeros(50, N, D)
x_seq = torch.cat((x_seq1,x_seq2), 0)
# 构建一层IF神经元
net = nn.Sequential(neuron.IFNode(v_threshold=9.0,
                                  v_reset=0.0,
                                  step_mode='s',
                                  store_v_seq=False))
print(net)

神经元数量N由输入维度(T,N)确定
使用脉冲神经元代替神经网络的激活函数

（2）根据输入按时间步更新神经元膜电位与输出
单步模式（默认）：逐步传播，深度优先遍历，内存占用小，适用于ANN2SNN
需要手动for循环按时间步计算

####################记录神经元状态####################
v_list = [] # 膜电位
s_list = [] # 神经元输出
####################单步模式：逐步传播####################
with torch.no_grad(): # 计算时关闭自动求导
    for i in range(T):
        y = net(x_seq[i])
        v_list.append(net[0].v)
        s_list.append(y)
functional.reset_net(net) # 重置神经元状态
####################可视化膜电位与输出####################
v_list = torch.cat(v_list).flatten()
s_list = torch.cat(s_list).flatten()
visualizing.plot_one_neuron_v_s(v_list.numpy(),
                                s_list.numpy(),
                                v_threshold=net[0].v_threshold,
                                v_reset=net[0].v_reset,
                                figsize=(12, 8),
                                dpi=100)

脉冲神经元是有状态的（self.v）
输入一批样本后需要进行复位：functional.reset_net(net)

（3）网络结构

Sequential(
  (0): IFNode(
    v_threshold=9.0, v_reset=0.0, detach_reset=False, step_mode=s, backend=torch
    (surrogate_function): Sigmoid(alpha=4.0, spiking=True)
  )
)

4、各时间步神经元的膜电位与输出

三、LIF神经元

1、神经元模型

（1）漏电整合发放(leaky integrate-and-fire)神经元：neuron.LIFNode

引入漏电项，无输入时膜电位恢复至静息电位，模拟离子扩散

（2）模型方程：
$I(t)-g(V(t)-E)=C\frac{dV(t)}{dt}$

I(t)：输入电流
V(t)：膜电位
C：膜电容
g：泄漏电导
E：静息电位

（3）阈下神经动力学方程：

$\left\{\begin{matrix} \begin{alignat*}{2} V[t-1]+X[t]-\frac{1}{\tau}(V[t-1]-V_{reset}),&\quad decay\_input=False\\ V[t-1]+\frac{1}{\tau}(X[t]-(V[t-1]-V_{reset})),&\quad decay\_input=True \end{alignat*} \end{matrix}\right.$

decay_input为False时，膜电位V的衰减由 $\frac{1}{\tau}(V-V_{reset})$ 控制
decay_input为True时，输入X[t]参与衰减，乘以系数 $\frac{1}{\tau}$

（4）充电方程

def neuronal_charge(self, x: torch.Tensor):
    if self.decay_input:
        if self.v_reset is None or self.v_reset == 0.:
            self.v = self.v + (x - self.v) / self.tau
        else:
            self.v = self.v + (x - (self.v - self.v_reset) / self.tau
    else:
        if self.v_reset is None or self.v_reset == 0.:
            self.v = self.v + x - self.v / self.tau
        else:
            self.v = self.v + x - (self.v - self.v_reset) / self.tau

（5）构建LIF神经元：layer = neuron.LIFNode()

构造参数：
①tau=2.0：膜电位时间常数
②decay_input=True：输入是否参与衰减
③v_threshold=1.0：阈值电压
④v_reset=0.0：重置电压
⑤surrogate_function=surrogate.Sigmoid()：梯度替代函数
⑥step_mode=‘s’：步进模式
⑦store_v_seq=False：是否保存所有时间步的膜电位self.v

2、神经元仿真

（1）构建输入与神经元层，前50步输入为1，后50步输入为0

####################构建输入####################
T = 100 # 时间步数
N = 1 # 样本数目
D = 1 # 输入维度/神经元数目
x_seq1 = torch.ones(50, N, D)
x_seq2 = torch.zeros(50, N, D)
x_seq = torch.cat((x_seq1,x_seq2), 0)
# 构建一层LIF神经元
net = nn.Sequential(neuron.LIFNode(tau=10.0,
                                    decay_input=True,
                                    v_threshold=0.9,
                                    v_reset=0.0,
                                    step_mode='m',
                                    store_v_seq=True))
print(net)

（2）根据输入按时间步更新神经元膜电位与输出
多步模式：逐层传播，广度优先遍历，并行速度更快，适用于梯度替代法
直接计算，不需要手动写for循环

####################监视器记录神经元状态####################
# 记录膜电位
monitor_v = monitor.AttributeMonitor('v_seq',
                                      pre_forward=False,
                                      net=net,
                                      instance=neuron.LIFNode)
# 记录输出
monitor_o = monitor.OutputMonitor(net=net,
                                  instance=neuron.LIFNode)
####################多步模式：逐层传播####################
with torch.no_grad(): # 计算时关闭自动求导
    net(x_seq)
functional.reset_net(net) # 重置神经元状态
####################可视化膜电位与输出####################
v_list = monitor_v.records[0].flatten()
s_list = monitor_o.records[0].flatten()
visualizing.plot_one_neuron_v_s(v_list.numpy(),
                                s_list.numpy(),
                                v_threshold=net[0].v_threshold,
                                v_reset=net[0].v_reset,
                                figsize=(12, 8),
                                dpi=100)

使用monitor监视器记录神经元输出与成员变量（膜电位）

（3）网络结构

Sequential(
  (0): LIFNode(
    v_threshold=0.9, v_reset=0.0, detach_reset=False, step_mode=m, backend=torch, tau=10.0
    (surrogate_function): Sigmoid(alpha=4.0, spiking=True)
  )
)

（4）各时间步神经元的膜电位与输出

输入为0时，膜电位自发泄漏

（5）当时间常数tau很大时，LIF神经元退化为IF神经元

net = nn.Sequential(neuron.LIFNode(tau=1.0e8,
                                    decay_input=False,
                                    v_threshold=9.0,
                                    v_reset=0.0,
                                    step_mode='m',
                                    store_v_seq=True))

Sequential(
  (0): LIFNode(
    v_threshold=9.0, v_reset=0.0, detach_reset=False, step_mode=m, backend=torch, tau=100000000.0
    (surrogate_function): Sigmoid(alpha=4.0, spiking=True)
  )
)

总结

不同神经元的主要差别在于阈下神经动力学方程（充电方程），而放电方程与重置方程相似；
IF神经元具有记忆效应，在无输入时，神经元膜电位会一直维持在当前状态；
LIF神经元引入漏电流项，在无输入时，膜电位会恢复至静息电位；
通过monitor和visualizing可以方便地监控和绘制神经元的膜电位与输出。

参考：
[1] SpikingJelly的文档——神经元
[2] 脉冲神经网络：模型、学习算法与应用
[3] 脉冲神经网络研究进展综述