Skip to content

Code for paper: Spiking Token Mixer: An Event-Driven Friendly Former Structure for Spiking Neural Networks

Notifications You must be signed in to change notification settings

brain-intelligence-lab/STMixer_demo

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

STMixer_demo

Code for paper: Spiking Token Mixer: An Event-Driven Friendly Former Structure for Spiking Neural Networks

Prerequisites

The Following Setup is tested and it is working:

  • Python>=3.5
  • Pytorch>=1.9.0
  • timm>=0.9.5
  • Cuda>=10.2

Description

  • We use Dspike surrogate gradient to realize the backward of step function.
  • LIF model is build in LIFSpike in models/layer.py.
  • You can use the following code to simply run this demo for cifar100:
 CUDA_VISIBLE_DEVICES=0,1  python -m torch.distributed.launch --nproc_per_node=2 train_timm.py --config data/cifar100.yml --model stmixerv3 --seed 40
  • Please change the relevant code on lines 834-836 of train_timm.py to change the network hyperparameter.

Pre-trained models

  • The STMixer_8_768 models we used on ImageNet are avilable here.

Citation

@inproceedings{
anonymous2024spiking,
title={Spiking Token Mixer:  A event-driven friendly Former structure for spiking neural networks},
author={Anonymous},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=iYcY7KAkSy}
}

About

Code for paper: Spiking Token Mixer: An Event-Driven Friendly Former Structure for Spiking Neural Networks

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published