如上图所示,目前已有的机器学习框架和脉冲网络仿真框架都无法彼此兼容,传统的脉冲网络框架无法与当下成熟的深度学习技术相结合。因此,SpikingJelly也就应运而生。同时期,国外也诞生了类似的Norse框架,两者达成了良性互动,共同推进了该领域的发展。
图源:Deeplearningincorporatingbiologicallyinspiredneuraldynamicsandin-memorycomputing
脉冲神经网络,简称SNN,被誉为第三代人工神经网络,是由大脑这样一个脉冲信号处理系统启发而构建的。大脑具有高级智能,并且功耗较低(有一种说法是仅相当于一个25w的灯泡)。通过借鉴大脑中的脉冲结构,SNN能够在保持低功耗的前提下,达到与ANN相当的性能。
基于SNN的诸多特性,其应用前景十分广阔,从简单的图像分类,动作识别,音频处理到复杂的音视频信号,强化学习,机器人控制任务,都是SNN施展自己本领的潜在舞台。
图源:TrueNorth:AcceleratingFromZeroto64MillionNeuronsin10Years
在SpikingJelly框架中,神经元的动态被描述为充电,放电,重置三个过程,与图中三种颜色分别对应。
框架提供了Neuron,Layer,Functional,Encoding四个基本模块:
框架还集成了多个神经形态数据集,因为大多数神经形态数据集需要专用软件读取,使用非常繁琐。因此SpikingJelly将常用神经形态数据集统一进行了封装,只需一行代码即可进行调用,其中事件数据被统一为(,,,)的Python字典格式,并提供将事件积分成帧数据的接口。来自其它框架,其他类型的数据也可以通过编码器间接使用。为了方便下载,在遵守原始协议或原作者同意的情况下,在OpenI上建立了数据集的下载镜像。
框架还提供了一些对神经元状态,脉冲特征进行可视化的API,方便研究者直接使用,不需要再去进行繁琐的绘图调整。
框架还提供了丰富的教程,通俗易懂,初学者和资深开发者均可快速上手,包含了方便初学者上手学习的基础教程,和实现具体任务的样例教程。
由于SNN的发放脉冲特征图中绝大部分位置都是0,针对SNN发放脉冲稀疏的特点,框架中加入了一些初步的工作,其中主要是稀疏矩阵运算的加速。根据实测,在矩阵规模大,且稀疏度高时,稀疏矩阵乘法的速度远高于普通矩阵乘法。
此外,SpikingJelly的模块在任务性能上也高度优化,不需要复杂的设计,只需要简单的网络和基本的神经元就可以达到STOA的性能。
在常用的benchmark上,框架实测ANN转SNN的性能损失很低,这使得很多不了解SNN的人也可以将他们的工作转换为相应的SNN,并达到接近的性能。
在神经元的实现中,隐状态被封装在模块内部,而PyTorch,Norse等框架的状态定义模块外部,这可以避免使用者手动维护神经元膜电位,并且在搭建多层网络时可以直接堆叠。
前面介绍的直接堆叠就是这里的逐步传播方式,支持用Sequential直接包装网络,使用简单。框架同样也提供了逐层传播的方式,并行度更高,但也会带来稍大的显存开销。两种方法在前向,反向传播时均完全等价。
下面是一个简单的网络示例,熟悉ANN的用户能够很快上手SNN的搭建。
ANN2SNN模块支持两种后端模型表示,同时也支持用户自定义对网络进行转换的方式。