DOCSEARCHbyAlgolia
Submit Docs

snnTorch Documentation

repository·master·Indexed 24 days ago

https://github.com/jeshraghian/snntorch

A Python package that extends PyTorch to enable gradient-based learning with spiking neural networks (SNNs), allowing spiking neurons to function as recurrent activation units. It includes implementations of neuron models such as Lapicque, the 0th Order Spike Response Model (snn.Alpha), and Synaptic Conductance-based LIF neurons, along with visualization tools like snntorch.spikeplot.raster and snntorch.spikeplot.animator.

Tokens
67.9K
Snippets
158
Records
259
Agent score
83%

What's inside snnTorch

  1. Overview of snnTorch components

    master

    snnTorch is a Python package for gradient-based learning with spiking neural networks (SNNs), extending PyTorch to support spiking neuron models as recurrent activation units. The library is organized into several functional modules:

    • snntorch: The core spiking neuron library, deeply integrated with PyTorch's autograd.
    • snntorch.spikegen: Tools for spike generation and converting data to spike formats.
    • snntorch.functional: Common arithmetic operations for spikes (e.g., loss functions, regularization).
    • snntorch.surrogate: Provides optional surrogate gradient functions for training.
    • snntorch.spikeplot: Visualization tools for spike-based data using matplotlib and celluloid.
    • snntorch.utils: Utility functions for datasets.
    • snntorch.import_nir / snntorch.export_nir: Enables importing/exporting to other SNN libraries via the NIR (Neural Intermediate Representation) framework.

  2. Regression with SNNs: Part I Overview

    master

    This tutorial series demonstrates how to perform regression using various spiking neuron models in snnTorch:

    • Part I: Trains the membrane potential of a Leaky Integrate-and-Fire (LIF) neuron to follow a given trajectory over time.
    • Part II: Uses LIF neurons with recurrent feedback for classification using regression-based loss functions.
    • Part III: Uses a complex Spiking LSTM network to train the firing time of a neuron.

  3. Use snntorch.spikeplot for animations and plots

    master
    The snntorch.spikeplot module is designed to reduce boilerplate code when generating various animations and plots for Spiking Neural Networks (SNNs). It is deeply integrated with matplotlib.pyplot and celluloid to facilitate easy visualization of spike trains and network activity.

  4. Available functional modules in snntorch

    master

    The snntorch.functional package is organized into several specialized submodules:

    • snntorch.functional.acc: Accuracy functions for evaluating spiking models.
    • snntorch.functional.loss: Loss functions designed for spiking neural networks (e.g., spike count-based losses).
    • snntorch.functional.reg: Regularization functions to constrain neuron behavior or weights.
    • snntorch.functional.quant: State quantization tools for simulating hardware constraints or discretized states.
    • snntorch.functional.probe: Probing mechanisms for analyzing internal spiking dynamics.

  5. Use snntorch.utils for dataset handling

    master
    The snntorch.utils module provides utility functions specifically designed for handling datasets within the snnTorch ecosystem. While the module is lightweight, it serves as a collection of helper functions to streamline data preparation and management for Spiking Neural Network (SNN) training workflows.

  6. Convert non-spiking data to spikes using snntorch.spikegen

    master

    The snntorch.spikegen module provides methods to convert tensors containing non-spiking data (such as continuous values) into discrete spikes. This is a common preprocessing step for Spiking Neural Networks (SNNs).

    Supported conversion methods include:

    • Rate coding: Encodes information in the frequency of spikes.
    • Latency coding: Encodes information in the timing of the first spike.
    • Delta modulation: Encodes information based on changes in the input signal.

  7. Exoplanet Hunter Tutorial Overview

    master

    The Exoplanet Hunter tutorial demonstrates how to train Spiking Neural Networks (SNNs) for time series data using light intensity measurements. Key learning objectives include:

    • Training SNNs for time series analysis.
    • Using SMOTE (Synthetic Minority Over-sampling Technique) to handle unbalanced datasets.
    • Evaluating model performance using metrics beyond accuracy, such as AUC, ROC, sensitivity, and specificity.
    • Applying SNNs to astronomy tasks where power efficiency is critical (e.g., deep space satellites).

  8. Define an SNN Model for IPU acceleration

    master

    When designing a model for IPU training, the loss function must be wrapped within the torch.nn.Module class. The forward method should return both the output (e.g., spike recordings) and the loss, specifically using poptorch.identity_loss to handle the loss return during training.

    Key implementation details:

    1. Initialize states: Use self.lif.init_leaky() to initialize membrane potentials.
    2. Temporal loop: Iterate through time steps to simulate the SNN.
    3. Loss wrapping: In the forward pass, if self.training is true, return spk2_rec, poptorch.identity_loss(self.loss_fn(mem2_rec, labels), "none").
    class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # ... layer definitions ...
        self.loss_fn = SF.ce_count_loss()

    def forward(self, x, labels=None):
        mem1 = self.lif1.init_leaky()
        mem2 = self.lif2.init_leaky()
        spk2_rec = []
        mem2_rec = []
       
        for step in range(num_steps):
            cur1 = self.fc1(x.view(batch_size,-1))
            spk1, mem1 = self.lif1(cur1, mem1)
            cur2 = self.fc2(spk1)
            spk2, mem2 = self.lif2(cur2, mem2)
            spk2_rec.append(spk2)
            mem2_rec.append(mem2)

        spk2_rec = torch.stack(spk2_rec)
        mem2_rec = torch.stack(mem2_rec)

        if self.training:
            return spk2_rec, poptorch.identity_loss(self.loss_fn(mem2_rec, labels), "none")
        return spk2_rec

  9. How surrogate gradients work in snnTorch

    master

    Because the discrete nature of spikes makes it difficult for torch.autograd to calculate analytical derivatives, snntorch uses surrogate gradients to approximate the backward pass. This allows gradient-based optimization to work with spiking neuron models.

    By default, snntorch overrides the default gradient using snntorch.surrogate.ATan. You can replace this with other approximations or probabilistic models available in the snntorch.surrogate module.

  10. Evaluate model performance using Sensitivity, Specificity, and AUC-ROC

    master

    In imbalanced datasets (like exoplanet detection), accuracy is often misleading. Use these metrics instead:

    • Sensitivity (Recall / True Positive Rate): $\frac{TP}{TP+FN}$. Measures the ability to correctly identify positive cases.
    • Specificity: $\frac{TN}{TN+FP}$. Measures the ability to correctly exclude negative cases.
    • AUC-ROC (Area Under the Receiver Operating Characteristic curve): Quantifies the ability to distinguish between classes.
      • Values $> 0.5$ (closer to 1) indicate good performance.
      • Values $\approx 0.5$ indicate random guessing.
      • Values $< 0.5$ indicate performance worse than random guessing.

    Note: In this specific tutorial context, increasing specificity often comes at the cost of sensitivity.

  11. Recast classification as a regression task

    master

    Instead of using Cross-Entropy loss (which drives correct classes to fire at all time steps and incorrect classes to not fire at all), you can recast classification as a regression task using Mean-Square Error (MSE).

    In this approach, you train the network to ensure the correct neuron fires a specific target number of times, while incorrect neurons fire a lower target number of times. This promotes sparser, more bio-inspired activity.

    For $n$ classes, the target $y$ and prediction $\hat{y}$ are vectors of length $n$, where $\hat{y}_i$ is the total spike count of the $i^{th}$ output neuron over the simulation runtime.

  12. When to use 2nd-order (Synaptic) vs 1st-order (Leaky) neurons

    master

    Choosing between snn.Synaptic (2nd-order) and snn.Leaky (1st-order) depends on your data and training requirements:

    Use 2nd-order (snn.Synaptic) when:

    • Long-term temporal relations: The input data has temporal dependencies occurring across long time-scales.
    • Sparse input patterns: The input spiking pattern is sparse.
    • Temporal coding: You need to control the precise timing of spikes. The synaptic current acts as a low-pass filter, smoothing the membrane potential and allowing for a finite rise time (introducing a delay between input and output spikes).

    Use 1st-order (snn.Leaky) when:

    • Simpler backpropagation: 1st-order models have fewer equations, making the gradient flow simpler.
    • General cases: For many simple datasets, optimal results often push the alpha parameter of a Synaptic model toward 0, effectively making it behave like a 1st-order model.