This repository is the official implementation of Exploring Loss Functions for Time-based Training Strategy in Spiking Neural Networks (NeurIPS 2023) [pdf].
Regarding our code, we have implemented two backends for neuron functions in our algorithm: The python backend and the cuda backend, where the cuda backend significantly accelerates the neuron functions. The cuda backend requires additional environment configuration as shown in below.
The environment can be configured by the following steps:
-
Install pytorch (torchvision, torchaudio)
-
Install corresponding version of cuda-nvcc if you want to use the cuda backend (you may also use the default nvcc if the version is right, or you have admin account and install the right version):
conda install -c nvidia cuda-nvcc
-
Install other packages through requirements.txt:
pip install -r requirements.txt
Modify the data path and network settings in the .yaml config files (in the ./networks folder).
We recommend you to run the code in Linux environment, since we use pytorch cuda functions in the backward stage and the compile process is inconvenient in Windows environment.
The backend option can be configured by setting backend: "cuda" or backend: "python" in the .yaml config files.
CUDA_VISIBLE_DEVICES=0 python main.py -config networks/config_file.yaml