Package: jaxsnn
jaxsnnjaxsnnis an event-based approach to machine-learning-inspired training and
simulation of SNNs, including support for neuromorphic backends (BrainScaleS-2).
We build uponjax, a Python library providing
autograd and XLA functionality for high-performance machine learning research.Building the SoftwareThe software builds upon existing libraries, such asjax,optax,
andtree-math.
When using the neuromorphic BrainScaleS-2 backend, the software stack of the
platform is required.We provide a container image (based on theSingularity format) including all build-time and runtime dependencies.
Feel free to download the most recent version fromhere.For all following steps, we assume that the most recent Singularity container is located at/containers/stable/latest.Github-based BuildTo build this project from public resources, adhere to the following guide:# 1) Most of the following steps will be executed within a singularity container#    To keep the steps clutter-free, we start by defining an aliasshopt-sexpand_aliasesaliasc="singularity exec --app dls /containers/stable/latest"# 2) Prepare a fresh workspace and change directory into itmkdirworkspace&&cdworkspace# 3) Fetch a current copy of the symwaf2ic build toolgitclonehttps://github.com/electronicvisions/waf-bsymwaf2icsymwaf2ic# 4) Build symwaf2iccmake-Csymwaf2ic
ln-ssymwaf2ic/waf# 5) Setup your workspace and clone all dependencies (--clone-depth=1 to skip history)c./wafsetup--repo-db-url=https://github.com/electronicvisions/projects--project=jaxsnn# 6) Load PPU cross-compiler toolchain (or build https://github.com/electronicvisions/oppulance)moduleloadppu-toolchain# 7) Build the project#    Adjust -j1 to your own needs, beware that high parallelism will increase memory consumption!c./wafconfigure
c./wafbuild-j1# 8) Install the project to ./bin and ./libc./wafinstall# 9) If you run programs outside waf, you'll need to add ./lib and ./bin to your path specificationsexportSINGULARITYENV_PREPEND_PATH=`pwd`/bin:$SINGULARITYENV_PREPEND_PATHexportSINGULARITYENV_LD_LIBRARY_PATH=`pwd`/lib:$SINGULARITYENV_LD_LIBRARY_PATHexportPYTHONPATH=`pwd`/lib:$PYTHONPATHStructurejaxsnnis split into two parts. Training ofSNNsis done in the init/apply style.Time Discretejaxsnn.discretesimulatesSNNsby treating time in a discrete way. It uses euler steps of a fixed size to advance the network forward in time which draws inspiration fromnorse.Time Continuousjaxsnn.eventtreats time continously and allows jumping from one event to the next one. It's core functionality consists of thestepfunction, which does three things:Find the next threshold crossingIntegrate the neuron to this point in timeApply the discontinuity after the treshold crossingjaxsnn.event.leaky_integrate_and_fireprovides multiple neuron types which can be used to build larger networks. Each neuron type defined the three functions mentioned above.BSS-2 Connectionjaxsnn.event.hardwareprovides functionality to connect to theBSS-2 systemand to conduct learning experimens on dedicated neuromorphic hardare.First StepsWe provide multiple examples for usage ofjaxsnn.Time discrete learning using surrogate gradients on the Yin-Yang dataset:python-mjaxsnn.discrete.tasks.yinyangEvent-based two layer feed-forward network with analytical gradients:python-mjaxsnn.event.tasks.yinyang_analyticalEvent-based recurrent network (with weights set up to emulate a two-layer feed-forward network) with gradients computed using the EventProp algorithm:python-mjaxsnn.event.tasks.yinyang_event_propBSS-2If you want to work with the BSS-2 system, a working example is provided:python-mjaxsnn.event.tasks.hardware.yinyangThe operation point calibration script issrc/pyjaxsnn/jaxsnn/event/hardware/calib/neuron_calib.py.
Example:srun-pcube--wafer69--fpga-without-aout0--ptycpython./neuron_calib.py\--waferW69F0\--threshold150\--tau-syn6e-6\--tau-mem12e-6\--refractory-time30e-6\--synapse-dac-bias1000--calib-dirsrc/pyjaxsnn/jaxsnn/event/hardware/calibIf you want to study the behaviour that different hardware artifacts (noise on the spike times) have on the performance of SNNs, check out this example:python-mjaxsnn.event.tasks.hardware.yinyang_mockYou can switch between an actual execution on BSS-2 and a pure software mock mode, in which the hardware is emulated by a second software network. You can
add noise to spikes from this first network or limit the dynamic range (like it is on BSS-2).DocsMultiple notebooks help you getting started withjaxsnn.event_based_snn.ipynbgives a great introduction on how to write event-based software for gradient-based learning with SNNs in JAXttfs.ipynbexplores how spikes times can be computed analytically and how a small network of LIF neurons can be constructedevent_prop.ipynbcompares the gradients of the EventProp algorithm analytical gradients (TTFS)TODONumeric: In theEventPropLIFneuron module, gradients currently do not flow correctly over multiple layers. This problem consists because the state of the input queue is not adjusted correctly in thecustom_vjp. It is therefore only possible to defined a multile layer networks via one recursive layer usingRecurrentEventPropLIFThe mapping between the hardware neuron modulesHardwareRecurrentLIF(which can simulate multiple feed-forward layers) and the populations / projections is not yet implemented cleanly and is hacked into the tasks (experiment returns a list of spikes for two layers, which are merged together, projections are hardcoded)Currently, in each task and experiment, small noise is added to the spike data from hardware. This is because thejaxsnngradient computation can not handle mutliple spikes withexactlythe same time, which can happen on BSS-2 because of the cycle resolution. This should either be moved to theexperimentclass directly, or the software should be adjusted to handle this case.Plotting: The plotting currently does not load from saved data, but runs at the end of each task. It should be set up to run stand-alone with data loaded from a file.AcknowledgementsThe software in this repository has been developed by staff and students
of Heidelberg University as part of the research carried out by the
Electronic Vision(s) group at the Kirchhoff-Institute for Physics.This work has received funding from the EC Horizon 2020 Framework Programme
under grant agreements 785907 (HBP SGA2) and 945539 (HBP SGA3), the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) under Germany's
Excellence Strategy EXC 2181/1-390900948 (the Heidelberg STRUCTURES Excellence
Cluster), the German Federal Ministry of Education and Research under grant
number 16ES1127 as part of the Pilotinnovationswettbewerb Energieeffizientes
KI-System, the Helmholtz Association Initiative and Networking Fund [Advanced
Computing Architectures (ACA)] under Project SO-092, as well as from the
Manfred Stärk Foundation, and the Lautenschläger-Forschungspreis 2018 for
Karlheinz Meier.LicensingSPDX-License-Identifier: LGPL-2.1-or-later