SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence

Abstract

Spiking neural networks (SNNs) aim to realize brain-inspired intelligence on neuromorphic chips with high energy efficiency by introducing neural dynamics and spike properties. As the emerging spiking deep learning paradigm attracts increasing interest, traditional programming frameworks cannot meet the demands of the automatic differentiation, parallel computation acceleration, and high integration of processing neuromorphic datasets and deployment. In this work, we present the SpikingJelly framework to address the aforementioned dilemma. We contribute a full-stack toolkit for preprocessing neuromorphic datasets, building deep SNNs, optimizing their parameters, and deploying SNNs on neuromorphic chips. Compared to existing methods, the training of deep SNNs can be accelerated 11×, and the superior extensibility and flexibility of SpikingJelly enable users to accelerate custom models at low costs through multilevel inheritance and semiautomatic code generation. SpikingJelly paves the way for synthesizing truly energy-efficient SNN-based machine intelligence systems, which will enrich the ecology of neuromorphic computing.

Motivation and introduction of the software framework SpikingJelly for spiking deep learning.

INTRODUCTION

Recently, artificial neural networks (ANNs), such as convolutional neural networks (CNNs) (1), recurrent neural networks (RNNs) (2), and transformers (3), have defeated most other methods and even surpassed the average ability levels of humans in some areas, including image classification (1, 4, 5), object detection (6–8), machine translation (3, 9–11), speech recognition (12, 13), and gaming (14, 15). These achievements are computer science–oriented because ANNs are mainly driven by gradient-based numerical optimization methods (16, 17), big data (18, 19), and massively parallel computing with graphics processing units (GPUs) (20, 21). Although neuroscience plays a diminished role in ANNs (22), insights from neuroscience are critical for building general human-level artificial intelligence (AI) systems (23, 24). The human brain is one of the most intelligent systems, having overwhelming advantages over any other artificial system in cognition and learning tasks such as transfer learning and continual learning (24). The neuroscientific community has been exploring biologically plausible computational paradigms to understand, mimic, and exploit the impressive feats of the human brain. Correspondingly, neuroscience-oriented spiking neural networks (SNNs) have been derived and are sometimes regarded as the next generation of neural networks (25). SNNs process and transmit information using short electrical pulses or spikes, making them similar to biological neural systems.

The neurons in SNNs are lower-level abstractions of biological neurons that collect signals from dendrites and process stimuli with nonlinear neuronal dynamics, which enable SNNs to be competitive candidates for processing spatiotemporal data (26, 27). Spike-based biological neural systems are extremely energy efficient, e.g., the human brain only consumes a power budget of approximately 20 W (28). Benefiting from event-driven computing on spikes, SNNs are up to 1000× more power efficient than ANNs (29) when running on tailored neuromorphic chips, including True North (30), Loihi (31), and Tianjic (29). The biological plausibility, spatiotemporal information processing capabilities, and event-driven computational paradigm of SNNs have attracted increasing research interest in recent years.

Emerging spiking deep learning methods

Because of their nondifferentiable spike trigger mechanisms and the complex spatiotemporal propagation processes, designing high-performance learning algorithms for SNNs is challenging. Traditional learning algorithms mainly incorporate biologically plausible unsupervised learning rules and primitive gradient-based supervised learning methods. Unsupervised learning algorithms inspired by the biological nervous system are applied in SNNs, including Hebbian learning (32), spike timing–dependent plasticity (STDP) (33), and their variants (34–38). Primitive supervised learning methods including SpikeProp (39), Tempotron (40), ReSume (41), and SPAN (42) achieve higher performance than biologically plausible unsupervised methods. However, these approaches are limited. Most SpikeProp-based methods only allow spiking neurons to fire no more than one spike, while Tempotron, ReSume, and SPAN cannot train SNNs with more than one layer. Thus, these primitive supervised learning methods can only solve tasks that are no harder than classifying the Modified National Institute of Standards and Technology (MNIST) dataset (43).

One of the key technologies that have led to the rapid progress of ANNs is deep learning (44), which optimizes the parameters of multilayer ANNs via backpropagation and learns high-dimensional representations of data with multiple levels of abstraction. To overcome the challenge of training SNNs, researchers have explored the application of deep learning methods to SNNs and achieved substantial performance improvements. Two of the most commonly used deep learning methods for SNNs are the surrogate gradient method (45–47) and the ANN-to-SNN conversion (ANN2SNN) (48–53). SNNs trained by the surrogate gradient method achieve high performance on complex datasets such as the Canadian Institute for Advanced Research (CIFAR) dataset (54), the Dynamic Vision Sensor (DVS) Gesture dataset (55) and the challenging ImageNet dataset (19) using only a few simulation time steps (56–61), while SNNs converted from ANNs attain almost the same accuracy as that of the original ANNs on the ImageNet dataset with dozens of simulation time steps (51, 62, 63). Because of the rapid progress achieved by deep learning methods, the applications of SNNs have been expanded beyond classification to other tasks including object detection (64–66), object segmentation (67, 68), depth estimation (69), and optical flow estimation (70). The boom exhibited by the research community indicates that spiking deep learning has become a promising research hotspot.

Demands for frameworks

Experience derived from the development of ANNs shows that software frameworks play a vital role in deep learning. Modern frameworks, including TensorFlow (71), Keras (72), and PyTorch (73), provide user-friendly frontend application programming interfaces (APIs) implemented by Python and high-performance backend accelerated by C++ libraries, e.g., Intel MKL and NVIDIA CUDA. Statistics show that the numbers of both new adopters and projects increase exponentially after the release of modern frameworks (74), indicating that these frameworks substantially reduce the workload required to build and train ANNs, help users quickly realize ideas, and make great contributions to the growth of deep learning research. The rapid development of machine learning frameworks has additionally accelerated the progress of the research community, which also highlights the importance of frameworks for spiking deep learning.

However, there is no mature framework available for spiking deep learning. Most existing frameworks for SNNs, including NEURON (75), NEST (76), Brian1/2 (77, 78), and GENESIS (79), can build detailed spiking neurons with complex neuronal dynamics, use numerical methods to approximate ordinary differential equations (ODEs), and simulate the biological neural system with high precision but do not integrate automatic differentiation, which is the core component required for gradient-based deep learning. These frameworks construct SNNs that are highly biologically plausible and can be used to investigate the functionality of real neural systems but are not designed to solve machine learning tasks. Nengo (80), SpykeTorch (81), and BindsNET (82) use simplified neurons with smaller numbers of ODEs than detailed neurons. With their low computational complexity brought by simplified neurons, these frameworks can implement some primary machine learning and reinforcement learning algorithms but still lack the modern deep learning capabilities of SNNs.

Because of the lack of available frameworks, researchers who want to combine advanced deep learning methods with SNNs have to build basic spiking neurons and synapses from scratch, resulting in repetitive and uncoordinated efforts. Deep SNNs involve a large number of matrix operations in both spatial and temporal dimensions of the data, which requires researchers to refine their codes to create high-performance programs that are accelerated by GPUs. Such workloads increase the burden on researchers. In the neuromorphic community, SNNs are frequently used to process data from neuromorphic sensors and deployed on neuromorphic computing chips, but data processing and deployment also require considerable time and effort. After skillful researchers implement their projects, inconsistent programming languages, coding styles, and model definitions generated by different authors will be difficult to reuse and divide the community. The efficiency of scientific research can be greatly improved if there exists a modern spiking deep learning framework that has at least the following three characteristics: exploits and accelerates spike-based operations; supports both simulations on CPUs/GPUs and deployment on neuromorphic chips; and provides a full-stack toolkit for building, training, and analyzing deep SNNs.

SpikingJelly: A modern framework for spiking deep learning

To solve the above issues and promote research on spiking deep learning, we present SpikingJelly, an open-source deep learning framework, to bridge deep learning and SNNs. Figure 1A shows a hierarchical overview of the architecture of SpikingJelly. On the basis of one of the most commonly used machine learning frameworks, PyTorch, SpikingJelly supports the simulation of SNNs on both CPUs and GPUs with autograd-enabled computation. Additional CUDA kernels are used for GPU-level acceleration beyond that provided by PyTorch. To achieve a balance between ease of use, flexible extensibility, and high performance, the deep learning aspects in SpikingJelly are elaborated into four sections: Components, Functions, Acceleration, and Networks (see the “Subpackages of deep learning” section in the Supplementary Materials). Components provide essential modules such as spiking neurons and synapses to build deep SNNs. Functions contain practical functions for training, simulating, analyzing, converting, quantizing, and deploying SNNs. Some modules in Functions are the functional formulations of the corresponding modules in Components. In this way, both object-oriented programming and procedural programming are supported to satisfy the diverse demands of users. Acceleration accelerates the SNN simulation process with extra semiautomatically generated CUDA kernels, thus exploiting the efficiency of low-level programming languages and reducing the development cost via the code generation technique. On the basis of the above subpackages, Networks provide classic and large-scale network structures such as Spiking ResNet for fast model reuse and primary SNN applications for beginners, which is favorably received by the community. Considering that neuromorphic datasets obtained from neuromorphic sensors (83–85) such as ATIS, DAVIS, and DVS are widely used in SNNs, SpikingJelly integrates neuromorphic dataset processing methods, including downloading, unifying data layouts, and reading interfaces with the general NumPy (86) format. With the quantize and exchange packages provided in Functions, compatibility with neuromorphic chips is also effectively implemented by providing low-bit quantizers for network weights and exchange functions for deploying SNNs.

Fig. 1. Overview of SpikingJelly.

(A) Architecture of the whole framework. (B) A code example of building and running an SNN, whose architecture is shown in (C and D). The execution times required for a single training iteration with T = 2, 4, 8, 16, and 32 and inference with T = 128 using Spiking ResNet-18 built by SpikingJelly, Norse, and snnTorch. (E) Comparisons of the ecological niche of SpikingJelly with those of other frameworks. (F) Typical adoptions from research based on SpikingJelly.

As a full-stack framework, SpikingJelly enables researchers to build SNNs with flexible and convenient APIs, simulate SNNs with extremely high efficiency, and deploy SNNs to edge AI devices. With SpikingJelly, a method for synthesizing a truly spike-based energy-efficient machine intelligence paradigm enriches the ecology of the research in this field.

RESULTS

Convenient and flexible SNN construction

SpikingJelly provides convenient and flexible APIs for users to build SNNs with a few steps. These APIs are designed in the acclaimed PyTorch style. Thus, those who are proficient in deep learning can also quickly start working with SpikingJelly. Figure 1B is an example of building and running an SNN, whose structure is shown in Fig. 1C. Note that we input a three-dimensional (3D) tensor X with a shape of (T, N, C), where T is the number of time steps, N is the batch size, and C is the number of features. Then, we use the multi_step_forward function to input X into the SNN, which implements a for-loop over time steps to send the inputs X[t] to the network and concatenate the outputs Y[t] to Y. In SpikingJelly, all hidden states are stored inside modules, which makes it convenient to build a network with sequential layers and access attributes from a specific module. Note that the hidden states need to be reset before sending a new sample, and we call the reset_net function in this example to reset the SNN.

High-performance simulation

SpikingJelly leverages the massively parallel computing power of GPUs to achieve extremely high SNN simulation performance. For quantitative assessment purposes, we evaluated the training and inference performance of modern spiking deep learning frameworks on the Spiking ResNet, which is the spiking version of the classic ResNet ANN structure (87) and is used frequently by SNN researchers (58, 60, 62, 88, 89). The compared frameworks were Norse (90) and snnTorch (91), which were proposed at almost the same time as SpikingJelly or afterward. The versions of Norse and snnTorch are 1.0.0 and 0.6.0, respectively, which are the latest versions. We used leaky-integrate-and-fire (LIF) neurons (92) in the network, as they are among the most commonly used spiking neurons in deep SNNs. The experiments are performed on an Ubuntu 18.04 server with an Intel Xeon Silver 4210R CPU, an NVIDIA A100-SXM-80GB GPU, and 256 gigabytes of memory.

Researchers usually simulate SNNs in two ways. The first approach is to train SNNs directly via the surrogate gradient method with backpropagation through time (BPTT). Limited by the fact that the memory consumption is approximately proportional to the number of time steps T, direct training often uses few time steps; e.g., T ≤ 32 for high-resolution datasets and T ≤ 64 for small-sized datasets. The second approach is to run SNNs with weights obtained via ANN2SNN. The ANN2SNN evaluation is restricted to inference, since this procedure does not require training. Most of the ANN2SNN methods are based on rate encoding and need far more time steps, e.g., T ≥ 128, than the surrogate gradient method. Conversion methods based on latency encoding rather than rate coding are also reported (93). Note that recent research has achieved a substantial reduction of T, such as surrogate learning methods (94) with T = 5 and ANN2SNN methods (51, 95) with T ≤ 64. Our evaluation involved two steps for the most common usage: (i) training with small time steps for surrogate learning and (ii) performing inference with large time steps for ANN2SNN.

First, we trained SNNs with T = 2, 4, 8, 16, and 32 using the surrogate gradient method. Figure 1D shows the execution times required by the three frameworks for a single training iteration on Spiking ResNet-18. It can be found that SpikingJelly has a notable training speed advantage over other frameworks, yielding higher speedup when T is larger, e.g., up to 11× when T = 32. Second, we tested SNNs in terms of inference with time steps of T = 128, 256, 512, and 1024 to simulate the usage of ANN2SNN. As the experimental results suggest, the inference times of the three frameworks are all proportional to the number of time steps T. We therefore report only the results obtained with T = 128 for simplicity. As Fig. 1D shows, SpikingJelly also achieves up to 2× speedup over other frameworks. The source codes and results of the experiments are provided in the Supplementary Materials.

Neuromorphic device support

Neuromorphic devices, including sensors and computing chips, are the key components of neuromorphic engineering and are also frequently used in spiking deep learning. SpikingJelly provides an efficient interface to incorporate SNNs with neuromorphic devices.

Most neuromorphic datasets are collected from neuromorphic sensors [e.g., N-MNIST (96) was collected from the ATIS sensor, and DVS Gesture was collected from the DVS128 camera], while others, such as ES-ImageNet (97), are converted from static images by software algorithms. Raw data collected through the sensor is stored in a specific format, e.g., AEDAT, which requires a specific binary decoding method. To reduce the cost of use, decoding methods for different formats are implemented in SpikingJelly. SpikingJelly provides both event-based and downsampled frame-based representations of datasets. The event representation is identical to the address event representation (AER) format, which is also the default neuromorphic interchip communication protocol. The i-th event is E[i] = (x_i, y_i, t_i, p_i), where (x_i, y_i) is the coordinate, t_i is the time stamp, and p_i is the polarity. Frame representations are widely used (45, 56, 57, 60, 98, 99) and are usually temporally downsampled from event representations. The frame F is a 4D tensor containing event counts with a shape of (T, C, H, W), where T is the number of frames, C is the number of channels, and H and W are the height and width of the frame, respectively. Prevalent event-to-frame downsampling methods are also provided by SpikingJelly.

The advantages of SNNs largely depend on the event-driven fashion, which can only be satisfied when neuromorphic computing chips are used. To achieve this goal, researchers first train SNNs on powerful CPUs/GPUs to obtain their optimized weights and then deploy the pretrained SNNs on neuromorphic computing chips for inference purposes. Such a pipeline is also supported by SpikingJelly. SpikingJelly provides the necessary conversion functionality to convert native SpikingJelly modules to the supported formats of neuromorphic computing chips. This allows native SpikingJelly SNNs to run on neuromorphic chips (see the “Exchange modules” section in the Supplementary Materials). Currently, two of the most commonly used neuromorphic computing chips are supported by SpikingJelly: Loihi from Intel and Lynxi KA200 from the Tianjic (29) group. Deployment on other neuromorphic computing chips can also be implemented by developing specific exchange modules.

Ecological niche

A wide variety of SNN frameworks with distinctive features and specific advantages for certain usages are available, which we denote as ecological niches of frameworks. To illustrate the unique ecological niche of SpikingJelly among others, we summarize the characteristics of commonly used SNN frameworks. In general, frameworks can be divided into three categories.

The first category contains classic biological frameworks including NEURON, NEST, and Brian2, which use the lowest-level abstractions of biological neuron models and integrate (or not integrate but easy to implement) biologically plausible learning rules such as STDP. To support GPUs, some classic frameworks provide subframeworks, e.g., CoreNEURON for NEURON and Brian2GENN for Brian2, to accelerate parts of the modules derived from the parent framework. Brian2 also provides Brian2Loihi to support the Loihi chip. With rapid development in recent decades, the classical frameworks have been widely used by neuroscientists and have grown into a large and active research community.

The second category, which can be considered as the intersection of neuroscience and computer science, includes Nengo and BindsNET. These frameworks are designed in the NumPy style and use neuron models of moderate complexity, which leads to a lower computational cost than classical frameworks. These frameworks are more compatible with GPUs because they provide GPU-supported backends such as OpenCL or are directly implemented on the basis of modern machine learning frameworks that fully support GPUs. More specifically, Nengo uses OpenCL-based NengoOCL or TensorFlow-based NengoDL to use GPUs, and BindsNET is based on PyTorch. In addition, Nengo supports Loihi, which is implemented by the NengoLoihi subframework. Biologically plausible rules are the main training algorithms in these frameworks. Note that Nengo also supports ANN2SNN via the NengoDL subframework. The rich ecosystem of Nengo makes it one of the most commonly used frameworks. In addition, the rising BindsNET has attracted thousands of followers within the GitHub community.

The third category considers the intersection of neuroscience and deep learning, which includes Norse, snnTorch, and SpikingJelly. All of these frameworks are based on PyTorch and support GPUs and automatic differentiation. High-level neuron model abstraction is used, allowing these frameworks to work easily with backpropagation. All these frameworks support at least one deep learning method. These frameworks have received a lot of attention from the research community due to the increasing interest in spiking deep learning in recent years. Compared to other frameworks, SpikingJelly has the advantage of full-stack integration, which supports neuromorphic datasets and chips, ANN2SNN, and surrogate gradients, as well as biologically plausible learning rules, and maximizes its simulation efficiency with specific optimization techniques for spike-based operations.

For a clear illustration purpose, we compare available SNN frameworks for SNNs in terms of five aspects: simulation performance (Performance), support for neuromorphic sensors and computing chips (Neuromorphic Support), community scale (Community), biological abstraction level (Biology), and support for surrogate learning and ANN2SNN (Deep Learning Support). The results are presented in Fig. 1E, which shows the ecological niche of SpikingJelly. For more details about the difference with frameworks, please refer to the “Comparison of SNN frameworks” section in the Supplementary Materials.

Adoptions by the community

The increasing number of adoptions by the community is symbolic of the success of a framework. Since being open-sourced in December 2019, SpikingJelly has been widely used in many spiking deep learning studies, including adversarial attack (100, 101), ANN2SNN (95, 102–106), attention mechanisms (107, 108), depth estimation from DVS data (69, 109), development of innovative materials (110), emotion recognition (111), energy estimation (112), event-based video reconstruction (113), fault diagnosis (114), hardware design (115–117), network structure improvements (60, 61, 118–121), spiking neuron improvements (56, 122–127), training method improvements (128–138), medical diagnosis (139, 140), network pruning (141–145), neural architecture search (146, 147), neuromorphic data augmentation (148), natural language processing (149), object detection/tracking for DVS/frame data (65, 66, 150), odor recognition (151), optical flow estimation with DVS data (152), reinforcement learning for controlling (153–155), and semantic communication (156). Figure 1F shows parts of these adoptions. Its widespread adoption by the community marks SpikingJelly as one of the most commonly used spiking deep learning frameworks.

Core modules of SpikingJelly

The key factor that enables SpikingJelly to achieve a combination of flexibility, efficiency, and completeness is the design philosophy of its modules. Some of the core components and functional modules are shown in Figs. 2 and 3, respectively.

Fig. 2. Component modules contained in SpikingJelly.

(A) The general discrete-time spiking neuron model described by Eqs. 5 to 7: the neuronal charge, fire, and reset equations. (B) Simulation of an LIF neuron. The input X[t] is 0 or 10 when 0 ≤ t < 128, 0 when 128 ≤ t < 196, and 3 when 196 ≤ t < 256. Here, {H^*[t]} records H[t] at all time steps and includes V[t] when S[t] = 1. (C) The inheritance relationships of an LIF neuron. (D) An example of the sigmoid surrogate function $\mathrm{\sigma }(x)=\frac{1}{1+\exp (-\mathrm{\alpha }x)}$, which is an approximator of the Heaviside function Θ(x). (E) The principle of the surrogate gradient method that uses Θ(x) during forward propagation to generate discrete binary spikes while using σ′(x) during backward propagation to obtain continuous float gradients. (F) The latency encoder that encodes larger values, which are shown by darker squares, for spikes with earlier firing times. (G) Illustration of the input image and the output spikes of the Poisson encoder. (H) Visualization of synthetic events with two polarities and (I) the frames downsampled from these events. (J) The workflow of dataset processing.

Fig. 3. Functional modules in SpikingJelly.

(A) Schematic plot demonstrating the use of a step-by-step or layer-by-layer propagation pattern to simulate a network with three layers and three time steps. X_i[t] is the input of the i-th layer at time step t, and Y_i[t] is the output of the i-th layer at time step t. $M_i^{\mathrm{ss}}$ denotes the i-th layer in the single-step mode, while $M_i^{\mathrm{ms}}$ denotes the i-th layer in the multistep mode. The complete computational graphs of the two propagation patterns are identical but constructed in different orders. (B) Functions of ANN2SNN that normalize weights and replace the rectified linear unit (ReLU) with an integrate-and-fire (IF) neuron layer. (C) The quantizer maps the input in the range [0,1] to the nearest fixed-point, the number of which is controlled by the bits. (D) The monitor works like a probe, recording the input, output, input gradients, output gradients, and attributes of the assigned module. (E) The STDP learner uses a monitor to record presynaptic spikes and presynaptic spikes, computes traces, and updates weights. (F) Diagram showing the process of slicing events from the DVS Gesture dataset to four bins and integrating them to four frames [shown in (G)].

Figure 2 (A to E) illustrates a spiking neuron and its components in SpikingJelly. As Fig. 2A shows, SpikingJelly uses three discrete-time equations, which are charge, fire, and reset equations, to describe the behaviors of spiking neurons. X[t], H[t], S[t], and V[t] are the input, membrane potential after charging, spikes, and membrane potential after resetting at time step t, respectively. Figure 2B shows the response of the LIF neuron in SpikingJelly for a given stimulus. Figure 2C shows the inheritance of the LIF neuron class. By inheriting the parent class, the LIF neuron can be easily implemented with only a few lines of code. Figure 2D shows the Heaviside function Θ(x), sigmoid surrogate function $\mathrm{\sigma }(x)=\frac{1}{1+\exp (-\mathrm{\alpha }x)}$ (where α is a hyperparameter for controlling the shape), and surrogate gradient σ′(x), which are key attributes of the spiking neuron shown in Fig. 2C. Figure 2E illustrates how surrogate learning is implemented in SpikingJelly. Θ(x) is applied during the forward propagation to generate spikes, and σ′(x) is applied during the backward propagation to calculate gradients. Figure 2 (F and G) visualizes two typical spiking encoders in SpikingJelly. The latency encoder shown in Fig. 2F encodes larger values (darker color) to earlier spikes. The Poisson encoder shown in Fig. 2G encodes the input value to spikes with a firing probability obeying the Poisson distribution. Figure 2H shows synthetic events with two polarities, and the frames downsampled by SpikingJelly are shown in Fig. 2I. The workflow used to process datasets in SpikingJelly is shown in Fig. 2J.

The modules in SpikingJelly have the step_mode attribute to decide whether to apply single-step or multistep forwarding. A module in the single-step mode receives X[t] and outputs Y[t], which are data at a single time step. In contrast, a module in the multistep mode receives X = {X[0], X[1], …, X[T − 1]} and outputs Y = {Y[0], Y[1], …, Y[T − 1]}, which are sequences with data at many time steps. Correspondingly, the SNN with all modules in a single-step or multistep mode follows the step-by-step or layer-by-layer propagation pattern, respectively, which is illustrated in Fig. 3A. The difference between these propagation patterns is the order in which the computational graphs are constructed. More specifically, the step-by-step propagation pattern is a depth-first-search (DFS), while the layer-by-layer pattern is a breadth-first-search (BFS). The propagation patterns in SpikingJelly are designed to satisfy specific user intentions (see the “Distinction of propagation patterns” section in the Supplementary Materials). The step-by-step approach is more flexible for building recurrent connections and is suitable for ANN2SNN because its memory consumption is not proportional to the number of time steps T during inference. The layer-by-layer pattern has the advantage of efficiency optimized by merging time steps into batch dimensions for stateless modules and fusing kernels for stateful modules in SpikingJelly. Figure 3B shows the ANN2SNN conversion functions, which normalize the weights and replace the rectified linear units (ReLUs) with integrate-and-fire (IF) neuron layers. Figure 3C shows the k-bit quantizer with k = 2 and 4, which quantizes the input x ∈ (0,1) to the nearest fixed-point value y. Quantizers support quantization-aware training, which quantizes weights during training. Note that the gradient of the quantizer is zero almost everywhere; hence, the surrogate gradient method is used. Figure 3D shows an example of using the monitor to record data. The monitor works like a probe, collecting data from prescribed modules. Inputs, outputs, input gradients, output gradients, and attributes can be recorded to meet the main data collection requirements. Figure 3E illustrates the STDP learner based on the trace method (157), which uses two monitors to record pre/postsynapse spikes. Then, the traces are updated, and the weights of the synapses are modified. Figure 3 (F and G) visualizes the event slicing and downsampling operations. In Fig. 3F, a sample from the DVS Gesture dataset is sliced into four bins, which are shown in four cuboids with different colors. Then, the events in each cuboid are accumulated, and four frames are generated in Fig. 3G. More details of modules in SpikingJelly can be found in Materials and Methods.

Typical applications

As a full-stack toolkit for spiking deep learning, SpikingJelly involves almost all SNN applications. For simplicity, we demonstrate three typical scenarios that use SpikingJelly in Fig. 4. The first case shown in Fig. 4A is an example of training a deep SNN to classify the neuromorphic DVS Gesture dataset. First, raw events from the DVS camera are preprocessed to tensors via the dataset subpackage in SpikingJelly. Then, we use the layer, neuron, and exchange subpackages to build and train a deep convolutional SNN for classifying the DVS Gesture dataset. The spikes flow through the convolutional layers, as shown in Fig. 4A. After training, we use the exchange subpackage to deploy the SNN on the neuromorphic Loihi chip for inference purposes. The second case shown in Fig. 4B is an example of using a deep SNN to solve the continuous control tasks from the OpenAI gym (158). First, the observation obtained from the state of the environment is encoded by different Gaussian receptive fields composed of population neurons (39), which effectively convert float values to spike trains. After processing, the discrete spike trains are converted to float values, which are the membrane potentials of subsequent layers of nonspiking neurons, to represent the action. Each module in the network can be implemented by the layer and neuron subpackages. The third case shown in Fig. 4C is an example of modeling the biological visual cortex by the deep SNN. First, we build and train the deep SNN on the ImageNet classification task by the layer and neuron subpackages. Then, to obtain the neural representations of the visual cortex and the model, the same visual stimuli are fed to the biological visual system and the pretrained SNNs. Last, the performance of SNNs in modeling the visual cortex can be quantitatively analyzed by calculating the resulting representation similarity. Source codes and experimental results for these applications are provided in the Supplementary Materials.

Fig. 4. Typical applications of SpikingJelly.

(A) The full-stack solution consists of preprocessing the raw events derived from the DVS Gesture dataset to tensors, building and training a deep SNN for classification, and deploying the SNN to the neuromorphic Loihi chip for inference. (B) Pipeline for reinforcement learning and control. The spike encoder transforms the observation obtained from the state of the environment in the OpenAI Gym into spike trains, and then, the spike trains are converted into continuous values in the spike decoder after SNN processing to represent the actions. (C) Workflow of measuring neural similarity between the biovisual system and SNNs. The same visual stimuli are fed to the biovisual system and the pretrained SNN to obtain neural representations of the two systems, and then, the brain-likeness of the model is assessed by measuring the representational similarity.

DISCUSSION

Although SNNs outperform ANNs in terms of biological plausibility and power efficiency, their applications were restricted to neuroscience rather than computational science because of the lack of available learning methods. With the introduction of deep learning methods, the performance of SNNs has been greatly improved, making spiking deep learning a new research hotspot. However, this emerging research area faces the dilemma that the classical software frameworks focus on neuroscience rather than deep learning, while new frameworks have not been developed. SpikingJelly is designed to satisfy the booming research interests of spiking deep learning (see the “Statistical trends” section in the Supplementary Materials).

Spiking deep learning is an emerging interdisciplinary field where researchers should be well versed in both neuroscience and deep learning. However, researchers who specialize in one research area may not be familiar with the other, which is consistent with our developers’ experience when answering questions and participating in discussions posed by users in GitHub. To mitigate the learning and the cost of use, SpikingJelly provides brief and convenient APIs. Classic models and frequently used training scripts are also included. With a few lines of code, users can easily build various types of SNNs and run their models even if they are not seasoned developers aware of the underlying implementation. This design philosophy of SpikingJelly frees users from painstaking coding operations when implementing more creative work.

Complex and diverse spiking neurons and synapses are the core components of SNNs. Modifying neurons and synapses by mimicking the biological neural system (56, 98, 141) or referring experience from deep ANNs (60, 89) is a practicable method for improving spiking deep learning. Researchers want to modify existing modules to define new classes of spiking modules by changing certain functions and properties. Researchers also expect that they should only need to write a few codes to make large changes to the behavior and performance of a model. Such a research paradigm is supported by SpikingJelly’s flexible APIs. Most of the modules in SpikingJelly are created by inheriting from their parents, overriding functions, and adding/deleting new attributes, which also provides a perfect reference for researchers to define new modules.

Deep learning typically uses datasets with massive numbers of samples and large-scale models (22). A larger number of training epochs is also used to achieve better performance (159). All the above characteristics are computationally intensive and hold for spiking deep learning. Moreover, because of the additional temporal dimension, deep SNNs have a higher level of computational complexity than deep ANNs. Thus, the simulation efficiency of deep SNNs is critical, especially for the recent research progress in terms of evaluating deep SNNs with more than 50 layers on the ImageNet dataset containing 1.28 million images has been a widely used performance baseline (58, 60, 62, 88, 89). Computational efficiency is emphasized in the design of SpikingJelly. The simulation process using SpikingJelly benefits from its infrastructure, PyTorch, which enables CPU acceleration with OpenMP/MKL and GPU acceleration with cuBLAS/cuDNN. Advanced acceleration with fusion operations is introduced by merging dimensions, semiautomatically generated CUDA kernels, and just-in-time (JIT) compiling, which brings extremely high training/inference efficiency. Equipped with these acceleration methods, SpikingJelly achieves state-of-the-art simulation speed, which frees researchers from wasting too much time waiting for lengthy deep SNN training processes.

Deep learning methods boost the performance of SNNs and make SNNs practical for use in real-life tasks (26). With the full-stack solution provided by SpikingJelly for building, training, and deploying SNNs, the boundaries of deep SNNs have been extended from toy dataset classification to applications with practical utility, including human-level performance classification, network deployment, and event data processing. Beyond the classic machine learning tasks, several frontier applications of deep SNNs have also been reported, including a spike-based neuromorphic perception system consisting of calibratable artificial sensory neurons (160), a neuromorphic computing model running on memristors (161), and the design of an event-driven SNN hardware accelerator (115). All the above evidence indicates that the advent of SpikingJelly will accelerate the boom of the spiking deep learning community. The future development scheme of SpikingJelly will keep track of the advances in neuromorphic computing, and the rapid symbiotic growth of SpikingJelly and the community will be witnessed.

MATERIALS AND METHODS

In this section, we first introduce the key modules of SpikingJelly and their design concepts. For clarity of exposition, we categorize modules into two categories: component modules and functional modules. Component modules are essential components of SNNs and are frequently used to define networks. Functional modules are modules or functions that simulate an SNN, accelerate its simulation process, record specific data, or modify variables. Then, the acceleration methods in SpikingJelly are discussed.

Component modules

Neuron model

Spiking neurons are the key components of SNNs. Neuroscientists prefer to use spiking neurons with complex neuronal dynamics, e.g., the Hodgkin-Huxley model (162) and the Izhikevich model (163), for high biological plausibility. When applying deep learning in SNNs, simplified spiking neuron models, including IF and LIF models, are more commonly used because of their straightforward neuronal dynamics and compact hyperparameters, which reduce the computational complexity and workload levels of researchers fine-tuning the training process of the network.

In general, the behavior of spiking neurons can be described by three processes: neuronal charging, neuronal firing, and neuronal resetting. Neuronal charging is also denoted as subthreshold dynamics, which can be defined as one or many ODEs. For example, the neuronal charging equation for an LIF neuron can be formulated as

$${\mathrm{\tau }}_m\frac{dV(t)}{dt}=-(V[t]-V_{\mathrm{rest}})+X(t)$$ (1)

where τ_m is the membrane time constant, V(t) is the membrane potential at time t, V_rest is the resting potential, and X(t) is the input current. The spiking neuron fires a spike S(t) = 1 when its membrane potential V(t) crosses the threshold V_th and remains silent [S(t) = 0] if V(t) < V_th, which can be described as

$$S(t)=\text{Θ}(V[t]-V_{\mathrm{th}})$$ (2)

where Θ(x) is the Heaviside step function defined by Θ(x) = 1 for x ≥ 0 and Θ(x) = 0 for x < 0. If the neuron fires a spike, then it will be reset. There are two types of resets, hard reset and soft reset. The hard reset, which resets V(t) to V_reset, is commonly used in directly trained SNNs because of its better experimental performance (164), while the soft reset, which subtracts V_th if the neuron fires, is adopted in ANN2SNN because of its lower theoretical fitting error for ReLU activations (50). The neuronal reset can be formulated as

$$\underset{\text{Δ}t\to 0^{+}}{lim}V(t+\text{Δ}t)=\{\begin{array}{c}{V}_{\mathrm{reset}},\text{hard reset}\\ V(t)-V_{\mathrm{th}},\text{soft reset}\end{array}$$ (3)

To simulate a spiking neuron, the continuous-time differential equations can be approximated by discrete-time differential equations. In practice, the simplest one-order Eulerian method is used because of its low computational cost. For example, the discrete-time version of Eq. 1 is

$${\mathrm{\tau }}_m(V[t]-V[t-1])=-(V[t-1]-V_{\mathrm{rest}})+X[t]$$ (4)

The neuronal charging equation is specific to different neurons, while the neuronal firing and neuronal resetting equations can be shared. Extending from our previous research (56), SpikingJelly uses three discrete-time equations to describe spiking neurons

$$H[t]=f(V[t-1],X[t])$$ (5)

$$S[t]=\text{Θ}(H[t]-V_{\mathrm{th}})$$ (6)

$$V[t]=\{\begin{array}{c}H[t]\cdot (1-S[t])+V_{\mathrm{reset}}\cdot S[t],\text{hard reset}\\ H[t]-V_{\mathrm{th}}\cdot S[t],\text{soft reset}\end{array}$$ (7)

To avoid confusion, we use H[t] to represent the membrane potential after charging but before firing and V[t] to represent the membrane potential after resetting. Equation 5 is the neuronal charging equation, and f is specific to a certain neuron. Equation 6 is the neuronal firing equation, and Eq. 7 is the neuronal reset equation. Figure 2A shows the general discrete-time neuron model described by these three equations.

Figure 2B shows the simulation of an LIF neuron built by SpikingJelly with τ_m = 50, V_th = 1, and V_reset = 0. The input X[t] is 0 or 10 when 0 ≤ t < 128, 0 when 128 ≤ t < 196, and 3 when 196 ≤ t < 256. Here, {H^*[t]} not only records H[t] at all time steps but also includes V[t′] for those time steps t′ when the neuron fires a spike (S[t′] = 1), which shows both fires and resets more clearly than only visualizing H[t] or V[t] alone.

Figure 2C takes an LIF neuron as the example and illustrates how SpikingJelly simplifies the module building process via multilevel inheritance. The StepModule is a module that can work in both single/multistep modes with a single tensor or a sequence input, and it is the base class of most modules in SpikingJelly (see the “Design of the step module” section in the Supplementary Materials). On the basis of the StepModule, the MemoryModule adds hidden states as extra attributes, making it suitable as the base class for stateful modules, including neurons and synapses in SNNs. Then, the general discrete-time neuron BaseNode can be easily designed by extending the MemoryModule with neuronal firing and resetting operations obeying Eqs. 6 and 7, respectively, and defining an empty neuronal charging function obeying Eq. 5 for child classes to complete. Last, we build the LIFNode inherited from the BaseNode by completing the neuronal charging function Eq. 5 and using an extra membrane potential parameter τ_m. Beyond the naive PyTorch implementation, SpikingJelly offers additional optional CUDA kernels for accelerating spiking neurons. The forward CUDA kernel LIFNodeFPTTKernel, the backward CUDA kernel LIFNodeBPTTKernel, the autograd function LIFNodeATGF, and the surrogate function of the LIF neuron are also extended through multilevel inheritance with several lines of code changes. Such a module design paradigm enhances the extensibility of SpikingJelly and reduces the workload required for researchers to develop new models. When researchers want to define a new kind of module, they only need to inherit the corresponding base class and complete or override some functions.

Surrogate gradients

The great success of deep learning models is largely due to their multiple processing layers trained by backpropagation and gradient descent (44). However, the derivative of the Heaviside function Θ(x) introduced by the spike trigger mechanism in Eq. 6 is +∞ at x = 0 and 0 at x ≠ 0, which blocks the gradient propagation process. To apply the gradient-based deep learning method in SNNs, the surrogate gradient method has been proposed (45, 46, 165, 166), which uses the derivative of the surrogate function σ(x) to define the derivative of Θ(x) during backpropagation. Commonly used surrogate functions are approximators of Θ(x), which have similar but smoother shapes; an example is the sigmoid surrogate function $\mathrm{\sigma }(x)=\frac{1}{1+\exp (-\mathrm{\alpha }x)}$ shown in Fig. 2D, where α is a hyperparameter used to control the shape.

As a surrogate function is used in the spike trigger mechanism, SpikingJelly packages it as a component of the spiking neuron. As shown in Fig. 2E, during forward propagation, the neuron uses S[t] = Θ(H[t] − V_th) to generate discrete binary spikes S[t] ∈ {0,1}. During backward propagation, the gradient is calculated by $\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }(H[t]-V_{\mathrm{th}})}=\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }S[t]}\cdot {\mathrm{\sigma }}^{\mathrm{\prime }}(H[t]-V_{\mathrm{th}})$, where ℒ is the loss.

Encoder

The spiking encoder is a unique module used in SNNs that encodes nonbinary input data into spikes. Figure 2F illustrates the latency encoder in SpikingJelly, which can be formulated as

$$S[t]=\{\begin{array}{c}1,t=t_f\\ 0,t\ne t_f\end{array}$$ (8)

$$t_f=[(T_{\max }-1)\cdot (1-x)]$$ (9)

where S[t] is the output spike at time step t, t_fis the firing time step, [·] is the rounding function, T_max is the latest firing time step, and x is the input value. The latency encoder encodes a larger x into an early firing time step t_f, indicating that a larger stimulus should cause a faster response, as has been observed in the retinal pathway (167).

The Poisson encoder is another commonly used encoder in SNNs that encodes inputs x ∈ [0,1] to Poisson events {s[t]} with

$$\begin{array}{r}P({\displaystyle \sum _{t=0}^{T-1}}s[t]=k)=\mathrm{Pr}(X=k)=\frac{T{x}^k\exp (-Tx)}{k!},k\in \{0,1,\dots ,T\}\end{array}$$ (10)

For simplicity, an approximate implementation is used, which is provided in SpikingJelly as

$$P(S[t]=1)=x$$ (11)

Figure 2G shows an example of an input image and the output spikes of the Poisson encoder. We find that the output spikes preserve most of the information from the original image. Other commonly used encoders, including weighted phase (168) and Gaussian tuning encoders (39), are also available in SpikingJelly.

Neuromorphic datasets

Neuromorphic datasets are frequently used as benchmarks for evaluating the temporal information processing abilities of SNNs (169). SpikingJelly provides commonly used datasets, including ASL-DVS (170), CIFAR10-DVS (171), DVS Gesture (55), ES-ImageNet (97), HARDVS (172), N-Caltech101 (96), N-MNIST (96), Nav Gesture (173), and Spiking Heidelberg Digits (174) (see the “Neuromorphic datasets” section in the Supplementary Materials). New datasets can be easily integrated by inheriting the dataset base and completing some processing functions.

Figure 2H visualizes a series of events with two polarities drawn in blue and green. The raw data collected from the sensor is stored in a specific format, e.g., AEDAT, which requires a specific binary decoding method. To reduce the cost of use, decoding methods for different datasets are implemented in SpikingJelly. Correspondingly, the outputs of the datasets in SpikingJelly are in the NumPy format, which is compatible with all Python applications. The number of events in a sample can be in the millions and can hardly be handled directly by the network. Widely used subsampling methods integrate events into frames. Figure 2I shows the four frames integrated from the events in Fig. 2H. SpikingJelly also provides frequently used integration methods.

Figure 2J shows the process flow of the datasets in SpikingJelly, which includes downloading from the source, extracting archives, decoding the raw data to an AER Python dictionary in the NumPy format, and downsampling events to frames. To add new datasets, the user only needs to inherit the base class and implement the “download”, “extract”, and “decode” functions, while another workload is handled by the base accelerated by multithreading, which also reveals the superior extensibility of SpikingJelly.

Functional modules

Step modes and propagation patterns

In SpikingJelly, each module has a variable attribute called step_mode, which controls whether the module uses the single-step mode or multistep mode (see the “Design of the step module” section in the Supplementary Materials). In the single-step mode, the module receives X[t] and outputs Y[t], which are data at a single time step. The input sequence for the 0-th layer is denoted as X₀ = {X₀[0], X₀[1], …, X₀[T − 1]}. When an SNN is stacked with L single-step modules {$M_0^{\mathrm{ss}}$, $M_1^{\mathrm{ss}}$, …, $M_{L-1}^{\mathrm{ss}}$}, the SNN is simulated in the step-by-step propagation pattern, which is shown in Algorithm 1A.

Algorithm 1: Propagation patterns.

(A) Propagation pattern: step by step

Require: an SNN stacked by L single-step modules $\{M_0^{\mathrm{ss}},M_1^{\mathrm{ss}},\dots ,M_{L-1}^{\mathrm{ss}}\}$, the input sequence for the 0-th layer X₀ = {X₀[0], X₀[1], …, X₀[T − 1]}.

Create an empty list Y_L−1 = {}

for t ← 0,1, …T − 1

for l ← 0,1, …L − 1

$X_{l+1}[t]=Y_l[t]=M_l^{\mathrm{ss}}(X_l[t])$

Append Y_L−1[t] in Y_L−1 = {Y_L−1[0], Y_L−1[1], …, Y_L−1[t − 1]}

Output Y_L−1 = {Y_L−1[0], Y_L−1[1], …, Y_L−1[T − 1]}

__________

(B) Propagation pattern: layer by layer

Require: an SNN stacked by L multistep modules $\{M_0^{\mathrm{ms}},M_1^{\mathrm{ms}},\dots ,M_{L-1}^{\mathrm{ms}}\}$,the input sequence for the 0-th layer X₀ = {X₀[0], X₀[1], …, X₀[T − 1]}

for l ← 0,1, …L − 1

$X_{l+1}=Y_l=M_l^{\mathrm{ms}}(X_l)$

Output Y_L−1 = {Y_L−1[0], Y_L−1[1], …, Y_L−1[T − 1]}

When switching the step_mode to the multistep mode, the module receives a sequence X = {X[0], X[1], …, X[T − 1]} and outputs a sequence Y = {Y[0], Y[1], …, Y[T − 1]}. Correspondingly, the SNN built with L multistep modules {$M_0^{\mathrm{ms}}$, $M_1^{\mathrm{ms}}$, …, $M_{L-1}^{\mathrm{ms}}$} can be simulated in a layer-by-layer manner, as Algorithm 1B shows.

Figure 3A shows how identical computational graphs are built in different orders by step-by-step and layer-by-layer propagation patterns. Considering that there are two dimensions, the time step and depth, in the computational graphs of SNNs, we can find that the step-by-step and layer-by-layer propagation patterns are DFS and BFS processes for traversing the computational graph, respectively. The network can switch between two propagation patterns easily by changing the step_mode attribute of each layer in SpikingJelly. The chosen propagation mode is determined by the user’s intent. The layer-by-layer pattern has the advantage of efficiency because the calculation over time steps can be executed in parallel for stateless layers and the fusion operations can be implemented for stateful layers, which will be discussed in the “Acceleration modules” section. However, when using the layer-by-layer pattern, the inputs X[t] at all time steps t must be given simultaneously, which is impossible when X[t + 1] depends on X[t], e.g., when we use the recurrent connections between layers. In such cases, the step-by-step pattern is preferred because of its flexible usage.

Conversion methods

In addition to surrogate gradient training, another way to obtain high-performance SNNs is through ANN-to-SNN conversion, also known as ANN2SNN. The conversion process is based on the fact that a rate-coded SNN contains an equivalence relation between the firing rate and the activation of the ANN (49, 95) under the conditions that the ANN should use ReLU activations and average pooling. On the basis of this fact, one can transform a trained ANN into a corresponding SNN.

To provide a convenient resolution to the ANN2SNN problem, SpikingJelly has a Converter module that normalizes the weight of the ANN and converts the ANN to an SNN by replacing the ReLUs with IF neuron layers, which is shown in Fig. 3B. The neuronal charging function of the IF neuron is

$$H[t]=V[t-1]+X[t]$$ (12)

which has not decayed and can approximate the ReLU in ANNs by firing rates. The implementation in SpikingJelly is based on the modeling of postsynaptic potential (51, 102, 103) so that the average postsynaptic potential is equal to the activation of the original ANN. We provide a VoltageHook module to record the maximum (or some percentile) of the hidden features as the scale of the postsynaptic current for the conversed SNN. The activation is then altered by the IF neuron layer to accomplish the conversion process. Furthermore, to perform fast inference (i.e., inference with few time steps), we provide a plug-and-play optimized membrane potential initialization (102) in SpikingJelly.

Quantizer

To perform inference with SNNs on resource-constrained hardware, including field-programmable gate arrays, neuromorphic chips, and mobile phones, network quantization is a necessary technique for reducing storage and computational costs. For example, a network quantized to 8 bits has a 4× reduction in its model size and memory bandwidth requirements and up to 4× throughput when compared with a 32-bit model. SpikingJelly provides a quantizer for quantization-aware training that can be used to quantize weights and neuronal dynamics during training. Figure 3C shows a typical k-bit quantizer in SpikingJelly, which maps an input x in the range [0,1] to the nearest fixed-point $y=q(x)=\frac{[(2^k-1)\cdot x]}{2^k-1}$, where [·] is the rounding function. Note that q′(x) is zero almost everywhere, and SpikingJelly also uses the surrogate method to redefine its gradients.

Monitor

The demands for recording data, e.g., the firing rates of spiking neuron layers, are frequently required by researchers. To satisfy the demand for recording data, SpikingJelly provides five general purpose monitors, which can monitor the input/output during forward/backward propagation and the attributes of specific layers. As Fig. 3D shows, a monitor acts like a probe, which is inserted into layers of the user-defined type and records the data with a custom transform λ. For example, if we are interested in the firing rates of all spiking neurons in an SNN, we can set the OutputMonitor with the layer type as the spiking neuron layer and the custom transform as $\mathrm{\lambda }(S)=\frac{1}{T}{\sum }_{t=0}^{T-1}S[t]$, where T is the number of time steps and S[t] is the spike at time step t. With this monitor, we can easily record the firing rate of each spiking neuron layer.

STDP learner

Combining the local unsupervised STDP learning process and the global supervised surrogate gradient has become a promising learning method for deep SNNs (175), and it can be implemented by the STDP learner in SpikingJelly. In the minimal example with two neurons shown in Fig. 3E, the STDP learner in SpikingJelly uses the “all-to-all” STDP with the trace method (157), in which all the pre- and postsynaptic spike pairs are considered. It uses monitors to record presynaptic spikes s_pre and postsynaptic spikes s_post. Then, the traces tr_pre[t], tr_post[t] are updated as

$$t{r}_{\mathrm{pre}}[t]=t{r}_{\mathrm{pre}}[t-1]-\frac{t{r}_{\mathrm{pre}}[t-1]}{{\mathrm{\tau }}_{\mathrm{pre}}}+s_{\mathrm{pre}}[t]$$ (13)

$$t{r}_{\mathrm{post}}[t]=t{r}_{\mathrm{post}}[t-1]-\frac{t{r}_{\mathrm{post}}[t-1]}{{\mathrm{\tau }}_{\mathrm{post}}}+s_{\mathrm{post}}[t]$$ (14)

where τ_pre and τ_post are the time constants of presynaptic and postsynaptic traces, respectively. Then, the weight w is updated by the traces as

$$\text{Δ}w[t]=F_{\mathrm{post}}(w[t])\cdot t{r}_{\mathrm{pre}}[t]\cdot s_{\mathrm{post}}[t]-F_{\mathrm{pre}}(w[t])\cdot t{r}_{\mathrm{post}}[t]\cdot s_{\mathrm{pre}}[t]$$ (15)

where F_pre and F_post are custom user-defined functions for controlling the amplitudes of synapse changes. In SpikingJelly, Δw can be added to the gradient $\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }w}$, indicating that the STDP learner can work together with the gradient descent method and various kinds of optimizers, including momentum stochastic gradient descent and adaptive moment estimation (Adam) (17).

Event downsampling methods

The event-to-frame downsampling method is widely used in deep SNNs with various kinds of slicing and integration methods. Considering these diverse usages, SpikingJelly supports both predefined and user-defined custom slicing and integration methods. A commonly used downsampling method is to integrate events with a fixed time duration ΔT, which can be formulated as

$$\mathit{F}[j,p,y,x]={\displaystyle \sum _{t_0+j\cdot \text{Δ}T\le t_i<t_0+(j+1)\cdot \text{Δ}T}}{I}_{p,x,y}(p_i,y_i,x_i)$$ (16)

where I_p,y,x(p_i, y_i, x_i) is an indicator function and equals 1 only when (p, y, x) = (p_i, y_i, x_i). Figure 3F visualizes the four slices obtained by the fixed time durations of events from a sample in the DVS Gesture dataset, and the four integrated frames are shown in Fig. 3G.

However, the fixed time duration-based integration generates frames with different frame numbers because the durations of the samples in neuromorphic datasets are not identical. Another frequently used method provided in SpikingJelly is fixed-frame number-based integration (56), which can be formulated as

$$F[j,p,y,x]={\displaystyle \sum _{i=j_l}^{j_r-1}}{I}_{p,x,y}(p_i,y_i,x_i)$$ (17)

where j_l and j_r are the indices created by a specific splitting method. For example, splitting by event number is performed as follows

$$j_l=\lfloor \frac{N}{T}\rfloor \cdot j$$ (18)

$$j_r=\{\begin{array}{c}\lfloor \frac{N}{T}\rfloor \cdot (j+1),j<T-1\\ N,j=T-1\end{array}$$ (19)

where ⌊·⌋ is the floor operation and N is the number of events. Splitting by event duration is performed as

$$\mathrm{\Delta }T=\lfloor \frac{t_{N-1}-t_0}{T}\rfloor$$ (20)

$$j_l=\underset{k}{\mathrm{argmin}}\{t_k\mid t_k\ge t_0+\text{Δ}T\cdot j\}$$ (21)

$$j_r=\{\begin{array}{c}\underset{k}{\mathrm{argmin}}\{t_k\mid t_k<t_0+\text{Δ}T\cdot (j+1)\}+1,j<T-1\\ N,j=T-1\end{array}$$ (22)

Note that fixed frame number-based integration can be used only once all events have arrived, which makes it unsuitable for a real-time system whose event durations are unpredictable. User-defined custom slicing and integration methods are also supported in SpikingJelly; they are implemented as callable functions and are used in the initialized args for the dataset class.

Acceleration nodules

Principle of acceleration

SpikingJelly accelerates the simulation of SNNs via two methods, which are the parallelization of sequential computations for stateless layers and the fusion of CUDA kernels for stateful layers. Stateless layers, including convolutional and linear layers, are widely used in SNNs. Similar to stateful layers, these layers also receive a sequence X = {X[0], X[1], …, X[T − 1]} with a shape of (T, N, …) as input. The difference between stateless and stateful layers is that the computation of Y[t] in stateless layers only depends on X[t]. Thus, the for-loop concerning the time steps is not necessary because the computation has no temporal dependence. SpikingJelly provides the SeqToANNContainer or its functional formulation seq_to_ann_forward to wrap stateless layers or their forward propagations. The SeqToANNContainer merges the time step dimension into the batch dimension by reshaping the input tensor from a shape of (T, N, …) to (T · N, …) before sending the input sequence to the stateless layer to execute the computation over all time steps in parallel and splits the outputs into a sequence Y = {Y[0], Y[1], …, Y[T − 1]} with the original sequence length. Figure 5A shows an example of using the SeqToANNContainer to wrap a 1D convolutional layer Conv1d, whose input has a shape of (N, C, L), where C is the number of channels and L is the input size. To avoid the for-loop concerning the time steps, using an (n + 1)-D convolution to implement the n-D convolution is also a practicable method and has been applied in some frameworks. The (n + 1)-D convolution method uses the time step dimension as the last dimension, and the input sequence has a shape of (N, C, …, T). Then, the n-D convolution imposed on the (n + 1)-D input sequence is implemented by the (n + 1)-D convolution with weight and stride values of 1 in the time step dimension. Figure 5B compares the execution times of different implementations that apply 2D convolution with 128 channels, a kernel size of 3, and a stride of 1 on the input with a shape of (T, N, C, H, W) = (T,16,128,64,64), which is a common convolutional operation for deep SNNs. In Fig. 5B, “SJ” is SpikingJelly’s method that merges the batch and time step dimensions, “RAW” is the plain for-loop conducted for the time steps, and “3D” is using the time step dimension as the depth dimension of the 3D convolution, which is the (n + 1)-D convolution method. The results show that SJ is the best method because its execution time increases much more slowly than those of the other two methods with increasing time steps.

Fig. 5. Acceleration in SpikingJelly.

(A) The use of seq_to_ann_forward for wrapping stateless layers and merging the time step dimension into the batch dimension to compute in parallel over time steps. (B) Comparison of execution times for different implementations of applying 2D convolution on sequential data. SJ is SpikingJelly’s method that merges the batch and time step dimensions. RAW is the plain for-loop for time steps, and 3D regards the time step dimension as the depth dimension of the 3D convolution. (C) Comparison between the execution times of the LIF neuron in SpikingJelly (SJ) accelerated by fused operations and the plain for-loop implemented by PyTorch (RAW). (D) Example of easy implementation of forward and backward kernels for an IF neuron with inheritance. (E) An example of the IF neuron that shows the complete workflow of CUDA neurons in SpikingJelly. The parameters in each neuron are marked with different colors. Correspondingly, the CUDA codes governed by these parameters are also highlighted with the same colors. Strikethroughs such as abc represent deletions of the code abc.

For stateful layers such as spiking neurons, Y[t] depends on not only X[t] but also the hidden states H[t − 1], which successively depend on X[t − 1], X[t − 2], …X[0]. Thus, the for-loop concerning time steps is inevitable. However, when using the naive for-loop implemented in PyTorch, one or a series of CUDA kernels are invoked at each time step. The frequent launching of many small CUDA kernels wastes time on calling overhead including memory access time and kernel launch time, which decreases the computational efficiency of the network. To accelerate the stateful layers, SpikingJelly fuses many small kernels with a few large kernels, which avoids the calling overhead. Figure 5C compares the execution times of the training and inference processes of the LIF neuron in multistep mode implemented by the CUDA kernel fusion method (SJ) with the plain for-loop implemented by PyTorch (RAW). The results show that SJ is much faster than RAW. It can also be found that as T increases, the execution times of both neurons increase linearly during the inference process. However, the execution time of RAW grows quadratically during training, while that of SJ still increases slowly with T. For more details about the effect of acceleration methods in SpikingJelly, refer to the “Ablation study of acceleration methods” section in the Supplementary Materials.

On the basis of the above two acceleration methods, it is recommended to train deep SNNs with the layer-by-layer propagation pattern for higher efficiency. As Algorithm 1B shows, each module receives sequence data, and then, the parallelization of the sequential computations for stateless layers and the fusion of CUDA kernels for stateful layers can be used. However, the spatial complexity of the layer-by-layer propagation pattern is 𝒪(T · N), even during inference. When T is large, e.g., hundreds and thousands of time steps in an ANN2SNN inference case, the layer-by-layer propagation pattern cannot be used because of memory restrictions. In such a case, the step-by-step propagation pattern can be used for inference because of its spatial complexity of 𝒪(N), which is not proportional to T.

Semiautomatic CUDA code generation

The fusion of operations into one CUDA kernel has the highest computational efficiency but involves a massive workload for developers. Standard practices include writing CUDA code, wrapping CUDA functions in C++ interfaces, and binding C++ to Python via pybind. To simplify the development process, SpikingJelly uses CuPy (176) to implement CUDA kernels. After defining a raw CUDA kernel in a Python string, CuPy wraps and compiles the string to make a CUDA binary library, which is cached and reused in subsequent runs. The introduction of CuPy avoids the trivialities of binding between CUDA, C++, and Python and helps users focus on a pure Python development environment.

A notable advantage of using CuPy is that CUDA codes are written in Python strings, indicating that we can easily control CUDA codes in a Python environment. Combined with the superior extensibility of SpikingJelly, this characteristic greatly simplifies the process of creating a CUDA kernel. As mentioned in Fig. 2C, the forward kernel and backward kernel for a specific neuron class are extended from the neuron kernel base class with a relatively small number of code changes (see the “Details of the neuron kernel” section in the Supplementary Materials). According to Eqs. 5 to 7 for forward propagation, the backward propagation process is defined as

$$\frac{d\mathcal{L}}{dH[t]}=\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }S[t]}\frac{dS[t]}{dH[t]}+(\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }V[t]}+\frac{d\mathcal{L}}{dH[t+1]}\frac{dH[t+1]}{dV[t]})\frac{dV[t]}{dH[t]}$$ (23)

$$\frac{d\mathcal{L}}{dX[t]}=\frac{d\mathcal{L}}{dH[t]}\frac{dH[t]}{dX[t]}$$ (24)

where $\frac{dS[t]}{dH[t]}$ is defined by the surrogate function, $\frac{dH[t+1]}{dV[t]},\frac{dH[t]}{dX[t]}$ is defined by Eq. 5, while the other parts are generic for different kinds of neurons. Thus, the neuron kernel base class in SpikingJelly is an incomplete CUDA kernel with vacant functions for the inheritor to complete, which includes Eq. 5 for the forward propagation process and $\frac{dH[t+1]}{dV[t]}$ and $\frac{dH[t]}{dX[t]}$ for the backward propagation process. For example, when we use an IF neuron whose neuronal charging function and backward propagation process are defined as

$$H[t]=V[t-1]+X[t]$$ (25)

$$\frac{dH[t+1]}{dV[t]}=1$$ (26)

$$\frac{dH[t]}{dX[t]}=1$$ (27)

then the forward propagation through time (FPTT) kernel IFNodeFPTTKernel and the backward propagation through time (BPTT) kernel IFNodeBPTTKernel are defined as shown in Fig. 5D. The FPTT kernel implements the neuronal dynamics of the IF neuron including Eqs. 12, 6, and 7 at T time steps in a single CUDA kernel. Correspondingly, the BPTT kernel implements the backward propagation of the neuronal dynamics at all time steps in a single CUDA kernel (see algorithm S3). We find that completing these empty functions requires only a few lines of code.

Figure 5E shows the complete workflow of CUDA neurons in SpikingJelly by taking an IF neuron as the example. Parameters in the neuron are marked with different colors. Correspondingly, the CUDA codes governed by these parameters are also highlighted with the same color. For example, when hard_reset is true, the Python class performs the following operations, which are also shown by the blue highlights in Fig. 5E: (i) Add float & v_reset in the forward and backward kernels’ arguments. (ii) Keep the hard reset commands v_v_seq[t + dt] = h_seq[t] * (1.0f - spike_seq[t]) + v_reset * spike_seq[t], and delete the sort reset commands v_v_seq[t + dt] = h_seq[t] - v_th * spike_seq[t] in the forward kernel. (iii) Keep the backward hard reset commands float grad_v_to_h = (1.0f) - spike_seq_t, and delete the backward soft reset commands float grad_v_to_h = (1.0f) in the backward kernel.

When detach_reset is true, S[t] in Eq. 7 is detached from the backward computation (177), and then, the corresponding CUDA codes are deleted, as the purple highlights in Fig. 5E show. The CUDA codes used by surrogate learning to define $\frac{dS[t]}{dH[t]}$ are generated by the surrogate function class, as the orange highlights in Fig. 5E show. After the CUDA code generation process, CuPy compiles these CUDA codes to the CUDA library, and forward/backward computation can be accelerated by fusion operations within the large CUDA kernel.

JIT acceleration

JIT is an alternative method for accelerating a network, which can compile Python codes into high-performance C++/CUDA programs while running. Some simple elementwise operations, e.g., the neuronal firing (Eq. 6) and neuronal resetting (Eq. 7) operations, can be wrapped into one CUDA kernel by JIT. As a common optimization method, JIT has not only lower performance but also a lower development cost than specific optimization methods, such as manually writing CUDA codes. Thus, SpikingJelly uses semiautomatically generated CUDA kernels for complex operations such as forward and backward kernels through time steps with gradients for spiking neurons and uses JIT to accelerate other simple functions such as the neuronal dynamics at one time step. A typical example is shown in Fig. 5C, in which the inference process of SJ is accelerated by JIT.

Supplementary Materials

This PDF file includes:

Supplementary Information

Figs. S1 to S28

Tables S1 and S2

References