Rockpool

Rockpool#

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.

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.

Contents