Homework 2 of AI3610.
# 1. Create a conda virtual enviornment.
conda create -n snn python=3.9 -y
conda activate snn
# 2. Install PyTorch >= 1.10 according to your CUDA version. For example:
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
# 3. Install the library
pip install -e .
The library is implemented with hydra, PyTorch and snnTorch.
The main script of training is scripts/train.py
and the main library is under snn_nmnist/
where snn_nmnist/logger.py
contains some functions of logging and snn_nmnist/utils.py
contains all necessary related util functions from model and dataloader building to result visualization, etc.
I support two input formats: static inputs and spike inputs. They are defined by the data.spike
flag in config. Default is false.
By default, the model is trained for
In this library, I implement a simple convolutional spiking neural network (CSNN) . The convolutional network architecture to be used is: 12C5-MP2-64C5-MP2-1024FC10
.
12C5
is a 5×5 convolutional kernel with 12 filtersMP2
is a 2×2 max-pooling function1024FC10
is a fully-connected layer that maps 1,024 neurons to 10 outputs
The detailed implementation can be found here. (See the build_model
function)
The system is implemented with hydra. The configuration can be easily figured out in cmd. I highly recommend users to run the script on a GPU.
# To use static data (default):
CUDA_VISIBLE_DEVICES=0 python scripts/train.py
# To use spike data:
CUDA_VISIBLE_DEVICES=0 python scripts/train.py exp_id=spike data.spike=true
All other flags can be found here.
The results and logs can be found in exp/{exp_id}/
after running my script, which contains the checkpoints, logs and some visualization results.
For static data, I achieve a final test accuracy of 99.12%. The log file can be found here. For spike data, I achieve a final test accuracy of 97.05%. The log file can be found here. You can reduce the data.subset
parameter or raise the train.num_epochs
parameter in configuration to obtain a better result.
Training loss curve and testing accuracy curve are shown as below:
Some spike counter examples:
Training loss curve and testing accuracy curve are shown as below:
Some spike counter examples:
I would like to thank the awesome snnTorch that I have referred to.
@article{eshraghian2021training,
title = {Training spiking neural networks using lessons from deep learning},
author = {Eshraghian, Jason K and Ward, Max and Neftci, Emre and Wang, Xinxin
and Lenz, Gregor and Dwivedi, Girish and Bennamoun, Mohammed and
Jeong, Doo Seok and Lu, Wei D},
journal = {arXiv preprint arXiv:2109.12894},
year = {2021}
}