惊蜇(SpikingJelly)
English | 中文
SpikingJelly 是一个基于 PyTorch ,使用脉冲神经网络(Spiking Neural Network, SNN)进行深度学习的框架。
SpikingJelly的文档使用中英双语编写: https://spikingjelly.readthedocs.io
安装
注意,SpikingJelly是基于PyTorch的,需要确保环境中已经安装了PyTorch,才能安装SpikingJelly。
从 PyPI 安装:
pip install spikingjelly
PyPI的安装包不包含CUDA扩展。如果想使用CUDA扩展,请从源代码安装:
通过GitHub:
git clone https://github.com/fangwei123456/spikingjelly.git
cd spikingjelly
python setup.py install
通过OpenI:
git clone https://git.openi.org.cn/OpenI/spikingjelly.git
cd spikingjelly
python setup.py install
以前所未有的简单方式搭建SNN
SpikingJelly非常易于使用。使用SpikingJelly搭建SNN,就像使用PyTorch搭建ANN一样简单:
class Net(nn.Module):
def __init__(self, tau=100.0, v_threshold=1.0, v_reset=0.0):
super().__init__()
# 网络结构,简单的双层全连接网络,每一层之后都是LIF神经元
self.fc = nn.Sequential(
nn.Flatten(),
nn.Linear(28 * 28, 14 * 14, bias=False),
neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset),
nn.Linear(14 * 14, 10, bias=False),
neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset)
)
def forward(self, x):
return self.fc(x)
这个简单的网络,使用泊松编码器,在MNIST的测试集上可以达到92%的正确率。 更多信息,参见时间驱动的教程。可以通过Python命令行直接运行这份代码,训练MNIST分类:
>>> import spikingjelly.clock_driven.examples.lif_fc_mnist as lif_fc_mnist
>>> lif_fc_mnist.main()
阅读spikingjelly.clock_driven.examples以探索更多先进的神经网络!
CUDA增强的神经元
SpikingJelly 提供了2种版本的神经元:用户友好的PyTorch版本,以及速度更快的CUDA版本。下图对比了各种类型的LIF神经元的运行时长:
设备支持
像使用PyTorch一样简单。
>>> net = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10, bias=False), neuron.LIFNode(tau=tau))
>>> net = net.to(device) # Can be CPU or CUDA devices
引用
如果您在自己的工作中用到了惊蜇(SpikingJelly),您可以按照下列格式进行引用:
@misc{SpikingJelly,
title = {SpikingJelly},
author = {Fang, Wei and Chen, Yanqi and Ding, Jianhao and Chen, Ding and Yu, Zhaofei and Zhou, Huihui and Tian, Yonghong and other contributors},
year = {2020},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/fangwei123456/spikingjelly}},
}
项目信息
北京大学信息科学技术学院数字媒体所媒体学习组 Multimedia Learning Group 和 鹏城实验室 是SpikingJelly的主要开发者。
开发人员名单可以在这里找到。
欢迎各位开发者参与此框架的开发!