Home

Awesome

Memory-Efficient Reversible Spiking Neural Networks

In this paper, we propose the reversible spiking neural network to reduce the memory cost of intermediate activations and membrane potentials during training of SNNs. We first extend the reversible architecture along temporal dimension and propose the reversible spiking block, which can reconstruct the computational graph of forward pass with a reverse process. On this basis, we present the RevSResNet and RevSFormer models, which are the reversible counterparts of the state-of-the-art SNNs.

<div align=center><img src="basicblocks.png" height=100% width=80%/> </div>

Main results on CIFAR10/CIFAR100

ModelTParam.CIFAR10 Top-1 AccCIFAR100 Top-1 Acc
RevSResNet21411.05M94.5375.46
RevSResNet37423.59M94.7776.34
RevSFormer-2-38445.76M95.2978.04
RevSFormer-4-38449.32M95.3479.04

Main results on CIFAR10-DVS/DVS128

ModelTFlops.CIFAR10 DVS Top-1 AccDVS 128 Top-1 Acc
RevSResNet24100.43G75.5094.44
RevSResNet24160.69G76.4096.53
RevSFormer-2-256103.78G81.4097.22
RevSFormer-2-256166.05G82.2097.57

Version of important requirements

timm==0.6.12; cupy==11.3.0; torch==1.10.2+cu113; spikingjelly==0.0.0.0.12;

Train

the training codes are put in each dataset folder in file train.sh

Contact

For help or issues using this git, please submit a GitHub issue.

For technical communications, please contact hongzhang99@zju.edu.cn.

Reference

If you find this repo useful, please consider citing:

@article{RevSNN,
  title={Memory-Efficient Reversible Spiking Neural Networks},
  author={Zhang, Hong and Zhang, Yu},
  journal={arXiv preprint arXiv:2312.07922},
  year={2023}
}

Our codes are based on PyTorch and SpikingJelly by Wei Fang.