snnTorch

snnTorch#

snnTorch is a deep learning simulator for spiking neural networks, built on PyTorch.

Supported Primitives in snntorch#

This library supports conversion of the following nodes to NIR:

  • Conv2d

  • Flatten

  • Affine

  • Linear

  • CubaLIF

  • LIF

  • AvgPool2d

This library supports conversion of the following nodes from NIR:

  • Conv2d

  • Flatten

  • Affine

  • Linear

  • CubaLIF

  • IF

  • LIF

  • AvgPool2d

Import a NIR graph to snnTorch#

from snntorch.import_nir import import_from_nir
import torch

import nir

# Create a NIR Network
affine_weights = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
affine_bias = torch.tensor([1.0, 2.0])
li_tau = torch.tensor([0.9, 0.8])
li_r = torch.tensor([1.0, 1.0])
li_v_leak = torch.tensor([0.0, 0.0])
li_v_thr = torch.tensor([1.0, 1.0])
li_v_reset = torch.tensor([0.0, 0.0])
nir_network = nir.NIRGraph.from_list(
    nir.Affine(affine_weights, affine_bias),
    nir.LIF(li_tau, li_r, li_v_leak, li_v_thr, li_v_reset)
)

# Import to snnTorch
snntorch_network = import_from_nir(nir_network)

Export a NIR graph from snnTorch#

import nir
import snntorch as snn
from snntorch.export_nir import export_to_nir
import torch

lif1 = snn.Leaky(beta=0.9, init_hidden=True)
lif2 = snn.Leaky(beta=0.9, init_hidden=True, output=True)

# Create a network
snntorch_network = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(784, 500),
    lif1,
    torch.nn.Linear(500, 10),
    lif2
)

sample_data = torch.randn(1, 784)

# Export to nir
nir_model = export_to_nir(snntorch_network, sample_data)

# Save to file
nir.write("nir_model.nir", nir_model)
Contents