This page was generated from docs/advanced/nir_export_import.ipynb. Interactive online version:

🧠 ⇌ 💾 Import / export between toolchains with NIR

BETA INFORMATION. SUBJECT TO API CHANGE.

Neuromorphic Intermediate Representation (NIR for short) is a new open-source package for interoperability between the new crop of SNN simulators. Rockpool supports import and export of (torch-backed) models using NIR.

These are supported using the nn.modules.to_nir() and nn.modules.from_nir() functions from the rockpool.nn.modules.torch package.

Using NIR requires the nir, nirtorch and torch dependencies. These can be installed alongside Rockpool with:

pip install 'rockpool[nir]'

This notebook shows you how to interact with NIR. Let’s start by generating a simple network.

Export from Rockpool to NIR

# - Import modules to build a simple network
from rockpool.nn.modules import LIFTorch, LinearTorch
from rockpool.nn.combinators import Sequential

try:
    from rich import print
except:
    pass

net = Sequential(
    LinearTorch((1, 2)),
    LIFTorch(2),
    LinearTorch((2, 4)),
    LIFTorch(4),
)
print(net)

TorchSequential  with shape (1, 4) {
    LinearTorch '0_LinearTorch' with shape (1, 2)
    LIFTorch '1_LIFTorch' with shape (2, 2)
    LinearTorch '2_LinearTorch' with shape (2, 4)
    LIFTorch '3_LIFTorch' with shape (4, 4)
}

To handle import and export, Rockpool provides the helper functions from_nir() and to_nir(). Let’s convert our Rockpool network to a NIR representation, and write it out to disk.

# - Import utilities to import and export NIR graphs
from rockpool.nn.modules import from_nir, to_nir
import nir
import torch

# - Convert the Rockpool network to the NIR representation
nir_graph = to_nir(net)
print(nir_graph)

# - Write the exported graph to disk
nir.write('rockpool.nir', nir_graph)

NIRGraph(
    nodes={
        'input': Input(input_type={'input': array([1, 1])}),
        '0_LinearTorch': Linear(weight=tensor([[-1.5277],
        [-2.2684]])),
        '1_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={}
        ),
        '2_LinearTorch': Linear(
            weight=tensor([[-0.4406,  0.1872],
        [ 0.2237, -1.0186],
        [ 0.2433,  0.9899],
        [-1.5922,  0.0632]])
        ),
        '3_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={}
        ),
        'output': Output(output_type={'output': array([1, 1, 4])})
    },
    edges=[
        ('input', '0_LinearTorch'),
        ('0_LinearTorch', '1_LIFTorch'),
        ('3_LIFTorch', 'output'),
        ('2_LinearTorch', '3_LIFTorch'),
        ('1_LIFTorch', '2_LinearTorch')
    ],
    input_type={'input': {'input': array([1, 1])}},
    output_type={'output': {'output': array([1, 1, 4])}},
    metadata={}
)

Import from NIR to Rockpool and deploy

Importing from NIR to Rockpool requires the from_nir() helper function, along with nir.read.

# - Read a NIR file from disk
nir_graph = nir.read('rockpool.nir')

# - Use `from_nir` to convert to Rockpool
torch_net = from_nir(nir_graph)

# - Alternatively, load and convert directly
torch_net = from_nir('rockpool.nir')
print(torch_net)

GraphExecutor(
  (0_LinearTorch): LinearTorch()
  (1_LIFTorch): LIFTorch()
  (2_LinearTorch): LinearTorch()
  (3_LIFTorch): LIFTorch()
  (input): Identity()
  (output): Identity()
)

Note that torch_net is a torch.nn.Module, and not a Rockpool-native module. This is a limitation caused by relying on the nirtorch package for NIR import. Nevertheless, this object contains Rockpool modules LinearTorch and LIFTorch, and can be simulated. Since it is a nirtorch.GraphExecutor module, it returns a tuple (output, state).

# - Simulate the module
output, state = torch_net(torch.rand(1, 100, 1))

It also supports the Rockpool as_graph() method, required for deployment.

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

# - Map, quantize and get a Xylo HW configuration
spec = mapper(torch_net.as_graph())
config, is_valid, msg = config_from_specification(**channel_quantize(**spec))

# - Produce a Xylo simulation matching the NIR graph
XyloSim.from_config(config)

XyloSim  with shape (16, 1000, 8)

See also

For more information about NIR, as well as how NIR integrates with other libraries, see https://neuroir.org.