Rockpool - NIR documentation

Rockpool (https://rockpool.ai) is a deep learning toolchain for deep spiking neural networks and other dynamical NN architectures, supporting deployment to Neuromorphic computing hardware.

Rockpool provides a convenient interface for designing, training and evaluating networks, which can operate both with continuous-time dynamics and event-driven dynamics.

You can import a model from NIR with the from_nir method in rockpool.nn.modules.torch.nir. Once it is in Rockpool, you can deploy your model to Synsense’s Xylo chips.

Supported Primitives in rockpool¶

This library supports conversion of the following nodes to NIR:

Affine
Linear
CubaLIF
LI
LIF

This library supports conversion of the following nodes from NIR:

Affine
Linear
CubaLIF
LI
LIF

Import a NIR graph to Rockpool¶

from rockpool.nn.modules import from_nir

import nir
import torch

# Supress warnings
import warnings
# warnings.filterwarnings('ignore')

# Create a NIR graph
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])
nir_network = nir.NIRGraph.from_list(
    nir.Affine(affine_weights, affine_bias), nir.LI(li_tau, li_r, li_v_leak)
)

# Create Rockpool model from NIR graph.
with warnings.catch_warnings():
    # Supress unrelated warnings from Rockpool
    warnings.simplefilter("ignore")
    rockpool_model = from_nir(nir_network)
print(rockpool_model)

GraphExecutor(
  (input): Identity()
  (affine): LinearTorch()
  (li): ExpSynTorch()
  (output): Identity()
)

Export a NIR graph from Rockpool¶

from rockpool.nn.modules import to_nir, LinearTorch, LIFTorch
from rockpool.nn.combinators import Sequential

n_in = 2
n_hidden = 4
n_out = 2
dt = 1e-3

# Create Rockpool model
net = Sequential(
    LinearTorch((n_in, n_hidden), has_bias=False),
    LIFTorch(n_hidden, dt=dt),
    LinearTorch((n_hidden, n_out), has_bias=False),
    LIFTorch(n_out, dt=dt),
)

# Convert model to NIR graph with a random input of representative shape
nir_graph = to_nir(net, torch.randn(1, 2))
print(nir_graph)

# Reload model from NIR
rockpool_model = from_nir(nir_graph)
print(rockpool_model)

NIRGraph(nodes={'input': Input(input_type={'input': array([1, 2])}), '0_LinearTorch': Linear(weight=tensor([[-0.5091,  0.9906],
        [ 0.6462,  0.4613],
        [-1.6254, -1.5708],
        [ 1.3620, -0.1414]])), '1_LIFTorch': CubaLIF(tau_syn=array([0.02, 0.02, 0.02, 0.02], dtype=float32), tau_mem=array([0.02, 0.02, 0.02, 0.02], dtype=float32), r=array([19.024588, 19.024588, 19.024588, 19.024588], dtype=float32), v_leak=array([0., 0., 0., 0.], dtype=float32), v_threshold=array([1., 1., 1., 1.], dtype=float32), w_in=array([1., 1., 1., 1.], dtype=float32), input_type={'input': array([4])}, output_type={'output': array([4])}, metadata={}), '2_LinearTorch': Linear(weight=tensor([[ 0.6191,  0.3268, -0.3174, -0.8611],
        [ 0.7899,  1.0755, -1.0221,  1.0548]])), '3_LIFTorch': CubaLIF(tau_syn=array([0.02, 0.02], dtype=float32), tau_mem=array([0.02, 0.02], dtype=float32), r=array([19.024588, 19.024588], dtype=float32), v_leak=array([0., 0.], dtype=float32), v_threshold=array([1., 1.], dtype=float32), w_in=array([1., 1.], dtype=float32), input_type={'input': array([2])}, output_type={'output': array([2])}, metadata={}), 'output': Output(output_type={'output': array([1, 1, 2])})}, edges=[('input', '0_LinearTorch'), ('1_LIFTorch', '2_LinearTorch'), ('2_LinearTorch', '3_LIFTorch'), ('3_LIFTorch', 'output'), ('0_LinearTorch', '1_LIFTorch')], input_type={'input': {'input': array([1, 2])}}, output_type={'output': {'output': array([1, 1, 2])}}, metadata={})
GraphExecutor(
  (input): Identity()
  (0_LinearTorch): LinearTorch()
  (1_LIFTorch): LIFTorch()
  (2_LinearTorch): LinearTorch()
  (3_LIFTorch): LIFTorch()
  (output): Identity()
)

/Users/dylan/SynSense Dropbox/Dylan Muir/LiveSync/Development/Projects/nir/rockpool_nir/rockpool/nn/modules/torch/torch_module.py:258: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  super().register_parameter(key, nn.Parameter(torch.tensor(value.data)))
/Users/dylan/SynSense Dropbox/Dylan Muir/LiveSync/Development/Projects/nir/rockpool_nir/rockpool/nn/modules/torch/linear_torch.py:111: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  torch.tensor(weight) if weight is not None else None,
/Users/dylan/SynSense Dropbox/Dylan Muir/LiveSync/Development/Projects/nir/rockpool_nir/rockpool/nn/modules/torch/nir.py:115: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  dt=torch.min(torch.tensor(_to_tensor(node.tau_mem / (1 + node.r)))).item(),

Deployment¶

As all Rockpool modules, imported networks support the .as_graph() method needed for mapping and deployment to neuromoprhic hardware. Note that Rockpool and several sub-dependencies are needed to support deployment. You can install these with

pip install 'rockpool[xylo]'

For more information, see the Rockpool tutorial notebooks describing deployment to Xylo: https://rockpool.ai/devices/xylo-overview.html

from rockpool.devices.xylo.syns61201 import mapper, config_from_specification, XyloSim
from rockpool.transform.quantize_methods import channel_quantize

spec = mapper(rockpool_model.as_graph())
config, is_valid, msg = config_from_specification(**channel_quantize(**spec))
xylo = XyloSim.from_config(config)

Caveats¶

NIR support in Rockpool is currently dependent on the nirtorch helper package, and only supports torch backend networks.