DOCSEARCHbyAlgolia
Submit Docs

SpikingJelly

repository·master·Indexed 24 days ago

https://github.com/fangwei123456/spikingjelly

A PyTorch-native deep learning framework for Spiking Neural Networks (SNNs) version 2.0.0.dev1. It supports large-scale training, inference, ANN-to-SNN conversion, and deployment to neuromorphic hardware. The framework provides multiple backends including torch, cupy, and triton, and includes specialized tools for SNN modeling, mixed-precision training, and benchmarking for models such as GPT-2 and Qwen2.5.

Tokens
181.7K
Snippets
321
Records
624
Agent score
83%

What's inside spikingjelly

  1. SpikingJelly core capabilities overview

    master

    SpikingJelly provides a comprehensive suite for SNN development:

    • SNN Modeling: Activation-based components (spiking neurons, surrogate gradients, stateful/stateless modules) and predefined models.
    • Training Workflows: PyTorch-native flows, online-learning utilities, and ANN2SNN conversion.
    • Performance: Multiple backends (torch, cupy, triton), FlexSN for custom kernels, and mixed-precision training (e.g., fp8).
    • Scaling: Memory-efficient training via spike compression (memopt), distributed execution, and precision policy tools.
    • Datasets: Support for various neuromorphic and event-based datasets with preprocessing pipelines.
    • Analysis: Profiling for FLOPs, SynOps, memory access, and inference energy estimation.
    • Interchange & Deployment: Interfaces for NIR, Lava, and Lynxi-oriented neuromorphic workflows.

  2. Transformer ANN2SNN Conversion Paths

    master

    SpikingJelly provides four distinct conversion paths in spikingjelly.activation_based.ann2snn for transforming Transformer architectures into Spiking Neural Networks (SNNs) or TD-equivalent models.

    1. TransformerTDEquivalentRecipe (Baseline): A direct Time-Difference (TD) equivalent operator replacement. It replaces Transformer operators with TD-equivalent modules that output floating-point differences rather than binary spikes. It does not perform calibration or maintain internal time loops via spikes.
    2. STATransformerRecipe (Spatio-Temporal Approximation): A model-level enhancement of the baseline. It uses cumulative differences and explicit step-mode readout, incorporating dataloader calibration and a spike encoder. It is a training-free approximation process.
    3. SpikeZIPTFQANNRecipe (SpikeZIP QANN-to-SNN): A module-tree path designed for inputs that are already SpikeZIP-compatible Quantization-Aware Neural Networks (QANNs).
    4. Qwen2SNNRecipe (Qwen2 Offline Multi-step Conversion): Calibrates Hugging Face Qwen2 causal LMs and converts all decoder blocks using SpikingJelly TD operators and signed activation-aware IF neurons.

    Important Notes:

    • STA Conversion: Not a strictly fully spike-driven SNN. The output of mode="spiking_encoder" is a quantized value (integer spikes $\times$ calibration threshold). When membrane potential is negative, the integer spike count can also be negative, meaning the output is not a binary spike tensor.
    • SpikeZIP Path: Uses ST-BIF neurons which support signed ternary outputs, rather than strictly binary spikes.

  3. Explore spikingjelly.activation_based.triton_kernel components

    master

    The spikingjelly.activation_based.triton_kernel package provides several specialized modules for high-performance SNN execution using Triton:

    • Predefined Neuron Kernels: Optimized kernels for specific neuron models (see spikingjelly.activation_based.triton_kernel.neuron_kernel).
    • FlexSN Implementation: Specialized implementations for FlexSN (see spikingjelly.activation_based.triton_kernel.flexsn).
    • Torch-to-Triton Transpiler: Tools to convert Torch operations to Triton (see spikingjelly.activation_based.triton_kernel.torch2triton).
    • Spike Compressors: Modules for spike compression (see spikingjelly.activation_based.triton_kernel.compress).
    • Utilities: Helper functions and dummy implementations located in triton_utils and dummy.

  4. Use the spikingjelly.timing_based package for SNN simulation

    master

    The spikingjelly.timing_based package provides tools for Spiking Neural Network (SNN) simulation using a timing-based approach. It is organized into several functional modules:

    • Encoding: Tools for converting continuous or discrete data into spike trains.
    • Neuron: Implementations of various spiking neuron models (e.g., LIF neurons) that process spike trains over time.

    For practical implementation, refer to the examples provided in the spikingjelly.timing_based.examples module.

  6. Distributed training helpers in spikingjelly.activation_based.distributed

    master

    The spikingjelly.activation_based.distributed package provides experimental distributed-training tools specifically designed for multi-step Spiking Neural Networks (SNNs) within the spikingjelly.activation_based framework. It leverages torch.distributed, DTensor, tensor parallelism, and FSDP2 to enable large-scale SNN training.

    Key capabilities include:

    • Model Analysis: Identifying stateful modules and candidates for tensor parallelism.
    • Distributed Planning: Building structured plans for distribution based on model analysis, network topology, and training objectives.
    • Runtime Application: Applying distribution plans to models to create an SNNDistributedRuntime.
    • Low-level Configuration: Manual control over SNN distribution via SNNDistributedConfig or configure_snn_distributed.

  7. Use Attention Layers in SpikingJelly

    master

    SpikingJelly provides attention mechanism implementations designed for deep Spiking Neural Networks (SNNs). These layers are categorized into two main types:

    1. Attention layers for convolutional SNNs: Designed to enhance convolutional architectures.
    2. Attention layers for Spiking Transformers: Specifically built for Spiking Transformer architectures.

    For detailed implementation details and architectural guidance on Spiking Transformers, refer to the Spikformer tutorial.

  8. Explore SpikingJelly API modules

    master

    SpikingJelly provides several specialized modules for Spiking Neural Network (SNN) development. The API is organized into the following functional areas:

    • Activation-based (RNN-like) SNNs: Located in spikingjelly.activation_based. These are suitable for users looking for RNN-like spiking dynamics.
    • Timing-based SNNs: Located in spikingjelly.timing_based. These focus on the temporal aspects of spiking.
    • Neuromorphic Datasets: Located in spikingjelly.datasets. Provides tools and data loaders specifically for neuromorphic computing tasks.
    • Visualization: Located in spikingjelly.visualizing. Tools for inspecting and visualizing SNN behaviors.
    • Package-level Configurations: Located in spikingjelly.configure. Used to manage global settings for the library.

  9. Explore the spikingjelly.activation_based package components

    master

    The spikingjelly.activation_based package provides tools for Spiking Neural Network (SNN) research and development using activation-based methods. The package is organized into several functional modules:

    Core Components

    • base: Base classes and fundamental abstractions.
    • neuron: Spiking neuron models.
    • surrogate: Surrogate gradient functions for training SNNs.
    • encoding: Methods for converting data into spike trains.
    • layer: SNN layers (e.g., convolutional, fully connected).
    • rnn: Recurrent SNN architectures.

    Functions and Utilities

    • functional: Functional interfaces for SNN operations.
    • distributed: Utilities for distributed training.
    • memopt: Memory optimization tools.
    • monitor: Tools for monitoring SNN behavior.
    • op_counter: Operation counters for complexity analysis.
    • learning: Learning algorithms and training utilities.
    • quantize: Quantization tools for SNNs.

    Specialized Modules

    • ann2snn: Tools for converting Artificial Neural Networks (ANNs) to SNNs.
    • model: High-level model definitions.
    • examples: Reference implementations and usage examples.
    • cuda_kernel / triton_kernel: Backend acceleration via custom CUDA or Triton kernels.
    • nir_exchange, lava_exchange, lynxi_exchange: Framework exchange and hardware deployment tools for NIR, Lava, and Lynxi platforms.

  10. Estimate computational complexity with spikingjelly.activation_based.op_counter

    master

    The spikingjelly.activation_based.op_counter package provides tools to measure the computational cost and energy consumption of Spiking Neural Networks (SNNs) based on their activations. It includes counters for various metrics such as FLOPs, MACs, SynOps, and memory access, as well as energy profilers like NeuroMC and SpikeSim.

    Available modules include:

    • FLOP Counter: Measures Floating Point Operations.
    • Memory Access Counter: Tracks memory access patterns.
    • MAC / AC / SynOp Counters: Measures Multiply-Accumulate (MAC), Accumulate (AC), and Synaptic Operations (SynOp).
    • Compute-Only Energy Estimator: Estimates energy consumption based on computation.
    • NeuroMC Energy Profiler: Advanced energy profiling.
    • SpikeSim Event-Driven Energy Profiler: Event-driven energy profiling.

  11. What is a Tempotron neuron

    master

    The Tempotron is a spiking neuron inspired by the ANN Perceptron. It performs weighted summation of inputs and outputs a binary result (0 or 1) for binary classification.

    Key Characteristics:

    • Membrane Potential ($V(t)$): Defined by the sum of weighted Postsynaptic Potentials (PSPs) triggered by input spikes at times $t_i$, plus a reset potential $V_{reset}$.
    • PSP Kernel ($K(t - t_i)$): Describes how the membrane potential responds to an input spike. It typically features an instantaneous spike followed by exponential decay.
    • Classification Logic: A single Tempotron acts as a binary classifier. It outputs 1 if the maximum membrane potential during the simulation period ($V_{t_{max}}$) reaches or exceeds a threshold ($V_{threshold}$), and 0 otherwise.
    • Capacity: A single Tempotron is limited to one spike per period and performs binary classification. Multiple Tempotrons can be used together for multi-class tasks.

  12. What is the spikingjelly.activation_based.op_counter module?

    master

    The spikingjelly.activation_based.op_counter module is a runtime profiling toolkit designed to measure model-side runtime costs and estimate energy consumption.

    Unlike static analysis tools that estimate costs based on layer shapes, op_counter uses PyTorch dispatch and module tracking to observe a real execution of your model. This is critical for Spiking Neural Networks (SNNs) because runtime metrics (like SynOps or energy) depend heavily on input sparsity and the actual activity of spikes during a specific forward or backward pass.

    Key capabilities include:

    • Counting FLOPs, memory accesses, SynOps, MACs, and ACs.
    • Building higher-level energy estimators based on these counts.
    • Providing per-module scope counts or global summaries.