OST2: One-step Spiking Transformer with Spiking Gate Unit

This is the code implementation for the paper "OST2: One-step Spiking Transformer with Spiking Gate Unit".

Requirements

timm==0.5.4

cupy==10.6.0

pytorch==1.13.1

spikingjelly==0.0.0.0.12

pyyaml==6.0

data prepare: your should download corresponding dataset and modify the dateset path 'data_dir:' at each dataset '.yml' file.

Training on ImageNet

Setting hyper-parameters in imagenet.yml

cd imagenet
python -m torch.distributed.launch --nproc_per_node=4 train.py

Training on cifar10/cifar100

Setting hyper-parameters in cifar10.yml/cifar100.yml

cd cifar
python train.py

Training on cifar10DVS

Setting hyper-parameters in train.py

cd dvs_cifar10
python train.py

Training on DVS128Gesture

Setting hyper-parameters in train.py

cd dvs_128_gesture
python train.py