Efficient training and inference on GPU for sinabs spiking neural network simulator.
Project description
sinabs-exodus
Sinabs-exodus is a plugin to the sinabs spiking neural network library. It can provide massive speedups in training and inference on GPU.
The tool is based on EXODUS^1, a formulation of backpropagation-through-time with surrogate gradients, that allows for efficient parallelization. EXODUS stands for EXact calculation Of Derivatives as Update to SLAYER. It builds upon the SLAYER[^2] algorithm, but uses mathematically accurate gradients and tends to be more robust to surrogate gradient scaling, making training less prone to suffer from exploding or vanishin gradients.
Some of the code in this library is loosely based upon slayerPytorch, the python implementation of SLAYER.
Getting started
Prerequisites
EXODUS uses CUDA for efficient computation, so you will need a CUDA-capable GPU, and a working installation of CUDA.
If you have CUDA installed, you can use the command
$ nvcc -V
to see the installed version. The last line should say something like Build cuda_xx.x.....
, where x.xx is the version.
Note that
$ nvidia-smi
does not show you the installed CUDA version, but only the newest version your Nvidia driver is compatible with.
You should also make sure that you have a PyTorch installation that is compatible with your CUDA version. To verify this, open a python console and run
import torch
print(torch.__version__)
The part after the +
in the output is the CUDA version that PyTorch has been installed for and should match that of your system.
Installation
After cloning this repository, the package can simply be installed via pip.
This is a namespace package
meaning that once installed this will be sharing its namespace with sinabs
package.
$ pip install .
Do not install in editable (-e
) mode.
Usage
If you have used sinabs before, using EXODUS is straightforward, as the APIs are the same.
You just need to import the spiking or leaky layer classes that you want to speed up from sinabs.exodus.layers
instead of sinabs.layers
.
Supported classes are:
IAF
LIF
ExpLeak
For example, instead of
from sinabs.layers import IAF
iaf = IAF()
do
from sinabs.exodus.layers import IAF
iaf = IAF()
Frequent Issues
CUDA is not installed or version does not match that of torch
If during installation you get an error, such as
RuntimeError:
The detected CUDA version (...) mismatches the version that was used to compile
PyTorch (...). Please make sure to use the same CUDA versions.
or
OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
CUDA is either not installed properly on your system or the version does not match that of torch (see above).
If you do have the correct version installed and the error still comes up, try to make sure that the environment variables such as PATH
and LD_LIBRARY_PATH
contain references to the correct directories. Please refer to NVIDIA's installation instructions for more details on how to do this for your system.
License
Sinabs-exodus is published under AGPL v3.0. See the LICENSE file for details.
Footnotes
[^2]: Sumit Bam Shrestha and Garrick Orchard. "SLAYER: Spike Layer Error Reassignment in Time." In Advances in Neural Information Processing Systems, pp. 1417-1426. 2018.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
File details
Details for the file sinabs-exodus-1.0.0.tar.gz
.
File metadata
- Download URL: sinabs-exodus-1.0.0.tar.gz
- Upload date:
- Size: 50.2 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.1 CPython/3.8.10
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 2b73115b0c1bcf6d3689726b17241e3c9087ff6496dfce70aa67401b90224ccb |
|
MD5 | 0c787f113649cfe4778a1d4a96dde506 |
|
BLAKE2b-256 | 1fd858d4ed92fe8c70e106a2c3dfe8cd0f278da2eda9b5cdc902d8dc417b352b |