diff --git a/README.md b/README.md index d5e2a10..818c3ac 100644 --- a/README.md +++ b/README.md @@ -5,41 +5,39 @@ # NIR - Neuromorphic Intermediate Representation +[![Nature Communications Paper](https://zenodo.org/badge/DOI/10.1038/s41467-024-52259-9.svg)](https://doi.org/10.1038/s41467-024-52259-9) +[![PyPI - Downloads](https://img.shields.io/pypi/dm/nir?logo=pypi)](https://pypi.org/project/nir/) +[![GitHub Tag](https://img.shields.io/github/v/tag/neuromorphs/nir?logo=github)](https://github.com/neuromorphs/NIR/releases) +[![Discord](https://img.shields.io/discord/1209533869733453844?logo=discord)](https://discord.gg/JRMRGP9h3c) + NIR is a set of computational primitives, shared across different neuromorphic frameworks and technology stacks. **NIR is currently supported by 7 simulators and 4 hardware platforms**, allowing users to seamlessly move between any of these platforms. -The goal of NIR is to decouple the evolution of neuromorphic hardware and software, ultimately increasing the interoperability between platforms and improving accessibility to neuromorphic technologies. -## Installation -NIR is installable via [pip](https://pypi.org/) -```bash -pip install nir -``` +NIR is useful when you want to move a model from one platform to another, for instance from a simulator to a hardware platform. + +> Read more about NIR in our [documentation about NIR primitives](https://neuroir.org/docs/primitives.html) -Check your [local framework]([https://neuroir.org/docs](https://neuroir.org/docs/support.html)) for NIR support. +> See [which frameworks are currently supported by NIR](https://neuroir.org/docs/support.html). ## Usage -> Read more in our [documentation about NIR usage](https://neuroir.org/docs) +> Read more in our [documentation about NIR usage](https://neuroir.org/docs) and see more examples in our [examples section](https://neuroir.org/docs/examples) -To end-users, NIR is just a declarative format that sits between formats and will hopefully be as invisible as possible. -However, it is possible to export Python objects or NIR files. +NIR serves as a format between neuromorphic platforms and will be installed alongside your framework of choice. +Using NIR is typically a part of your favorite framework's workflow, but follows the same pattern when you want to move from a *source* to a *target* platform: ```python -import nir -# Write to file -nir.write("my_graph.nir", nir_graph) - -# Read file +# Define a model +my_model = ... +# Save the model (source platform) +nir.write("my_graph.nir", my_model) +# Load the model (target platform) imported_graph = nir.read("my_graph.nir") ``` -## About NIR -> Read more in our [documentation about NIR primitives](https://neuroir.org/docs/primitives.html) - -On top of popular primitives such as convolutional or fully connected/linear computations, we define additional compuational primitives that are specific to neuromorphic computing and hardware implementations thereof. -Computational units that are not specifically neuromorphic take inspiration from the Pytorch ecosystem in terms of naming and parameters (such as Conv2d that uses groups/strides). - +See our [example section](https://neuroir.org/docs/examples) for how to use NIR with your favorite framework. ## Frameworks that currently support NIR +> Read more in our [documentation about NIR support](https://neuroir.org/docs/support.html) | **Framework** | **Write to NIR** | **Read from NIR** | **Examples** | | --------------- | :--: | :--: | :------: | @@ -54,7 +52,7 @@ Computational units that are not specifically neuromorphic take inspiration from ## Acknowledgements -This work was originally conceived at the [Telluride Neuromorphic Workshop 2023](tellurideneuromorphic.org) by the authors below (in alphabetical order): +This work was originally conceived at the [Telluride Neuromorphic Workshop 2023](https://tellurideneuromorphic.org) by the authors below (in alphabetical order): * [Steven Abreu](https://github.com/stevenabreu7) * [Felix Bauer](https://github.com/bauerfe) * [Jason Eshraghian](https://github.com/jeshraghian) @@ -64,7 +62,7 @@ This work was originally conceived at the [Telluride Neuromorphic Workshop 2023] * [Sadique Sheik](https://github.com/sheiksadique) * [Peng Zhou](https://github.com/pengzhouzp) -If you use NIR in your work, please cite the [following arXiv preprint](https://arxiv.org/abs/2311.14641) +If you use NIR in your work, please cite the [following paper](https://www.nature.com/articles/s41467-024-52259-9) ``` article{NIR2024, diff --git a/docs/source/_config.yml b/docs/source/_config.yml index 3d06869..a705a71 100644 --- a/docs/source/_config.yml +++ b/docs/source/_config.yml @@ -10,7 +10,16 @@ repository: execute: execute_notebooks: off +parse: + myst_enable_extensions: + - amsmath + launch_buttons: notebook_interface: "jupyterlab" binderhub_url: "https://mybinder.org/v2/gh/neuromorphs/nir/main?urlpath=lab" - colab_url: "https://colab.research.google.com" \ No newline at end of file + colab_url: "https://colab.research.google.com" + +sphinx: + extra_extensions: + - 'sphinx.ext.autodoc' + \ No newline at end of file diff --git a/docs/source/_toc.yml b/docs/source/_toc.yml index 80178af..c9e4b3a 100644 --- a/docs/source/_toc.yml +++ b/docs/source/_toc.yml @@ -14,19 +14,24 @@ parts: - file: primitives - file: support title: Platform support -- caption: Examples - chapters: - - file: examples/lava/nir-conversion - - file: examples/nengo/nir-conversion - - file: examples/norse/nir-conversion - - file: examples/rockpool/nir-conversion - - file: examples/sinabs/nir-conversion - - file: examples/snntorch/nir-conversion - - file: examples/spinnaker2/import - - file: examples/spyx/conversion - - file: examples/snntorch_to_norse + - file: examples/index + sections: + - file: examples/lava/nir-conversion + - file: examples/nengo/nir-conversion + - file: examples/norse/nir-conversion + - file: examples/rockpool/nir-conversion + - file: examples/sinabs/nir-conversion + - file: examples/snntorch/nir-conversion + - file: examples/spinnaker2/import + - file: examples/spyx/conversion + - file: examples/snntorch_to_norse - caption: Developer guide chapters: - file: porting_nir - - file: api_design + - file: dev_pytorch + - file: dev_jax - file: contributing +- caption: API documentation + chapters: + - file: api_design + - file: doctrees \ No newline at end of file diff --git a/docs/source/contributing.md b/docs/source/contributing.md index 3a3a638..3b1c6d3 100644 --- a/docs/source/contributing.md +++ b/docs/source/contributing.md @@ -1,5 +1,9 @@ # Contributing +NIR is a community-led initiative, and we welcome contributions from everyone. +Here, we outline some technical details on getting started. +Join the conversation on our [Discord server](https://discord.gg/JRMRGP9h3c) or [GitHub](https://github.com/neuromophs/nir) if you have any questions. + ## Developer guide: Getting started Use the standard github workflow. diff --git a/docs/source/dev_jax.md b/docs/source/dev_jax.md new file mode 100644 index 0000000..0b91cc1 --- /dev/null +++ b/docs/source/dev_jax.md @@ -0,0 +1,6 @@ +# Developing JAX extensions + +JAX is a popular deep learning framework that more and more of the NIR-supported libraries are built on. +For PyTorch, we have built the [`nirtorch` package](https://github.com/neuromorphs/nirtorch), but *no such package exists for JAX*. +If you're interested in developing such a package, please reach out to us! +Either on [Discord](https://discord.gg/JRMRGP9h3c) or by [opening an issue](https://github.com/neuromorphs/NIR/issues). \ No newline at end of file diff --git a/docs/source/dev_pytorch.md b/docs/source/dev_pytorch.md new file mode 100644 index 0000000..883cce2 --- /dev/null +++ b/docs/source/dev_pytorch.md @@ -0,0 +1,69 @@ +# Developing PyTorch extensions + +PyTorch is a popular deep learning framework that many of the NIR-supported libraries are built on. +We have built the [`nirtorch` package](https://github.com/neuromorphs/nirtorch) to make it easier to develop PyTorch extensions for the NIR-supported libraries. +`nirtorch` helps you write PyTorch code that (1) exports NIR models from PyTorch and (2) imports NIR models into PyTorch. + +## Exporting NIR models from PyTorch +Exporting a NIR model requires two things: exporting the model's nodes and edges. + +### Exporting edges +Exporting edges is slightly complicated because PyTorch modules can have multiple inputs and outputs. +And because PyTorch modules are connected via function calls, which only happen at runtime. +Therefore, we need to trace the PyTorch module to get the edges with some sample input. +Luckily, `nirtorch` package helps you do exactly that. +It works behind the scenes, but you can read about it in the [`to_nir.py` file in `nirtorch`](https://github.com/neuromorphs/NIRTorch/blob/main/nirtorch/to_nir.py#L11). + +### Exporting nodes +The only thing we really have to do to use `nirtorch` is to export modules. +Since all PyTorch modules inherit from the `torch.nn.Module` class, exporting the nodes is straightforward: we simply need a function that looks at a PyTorch module and returns the corresponding NIR node. +Assume this is done in a function called `export_node`. + +```python +import nir +import torch + +class MyModule(torch.nn.Module): + weight: torch.Tensor + bias: torch.Tensor + + +def export_node(module: torch.nn.Module) -> Node: + # Export the module to a NIR node + if isinstance(module, MyModule): + return nir.Linear(module.weight, module.bias) + ... +``` +This example converts a custom Linear module to a NIR Linear node. + +### Putting it all together +The following code is a snippet taken from the [Norse library](https://github.com/norse/norse) that demonstrates how to export custom PyTorch models to a NIR using the `nirtorch` package. +Note that we only have to declare the `export_node` function for each custom module we want to export. +The edges are traced automatically by the `nirtorch` package. + +```python +def _extract_norse_module(module: torch.nn.Module) -> Optional[nir.NIRNode]: + if isinstance(module, LIFBoxCell): + return nir.LIF( + tau=module.p.tau_mem_inv, + v_th=module.p.v_th, + v_leak=module.p.v_leak, + r=torch.ones_like(module.p.v_leak), + ) + elif isinstance(module, torch.nn.Linear): + return nir.Linear(module.weight, module.bias) + elif ... + + return None + +def to_nir( + module: torch.nn.Module, sample_data: torch.Tensor, model_name: str = "norse" +) -> nir.NIRNode: + return extract_nir_graph( + module, _extract_norse_module, sample_data, model_name=model_name + ) +``` + +## Importing NIR models into PyTorch +Importing NIR models into PyTorch with `nirtorch` is also straightforward. +Assuming you have a NIR graph in the Python object `nir_graph` (see [Usage](#usage)) \ No newline at end of file diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md new file mode 100644 index 0000000..5568f19 --- /dev/null +++ b/docs/source/examples/index.md @@ -0,0 +1,33 @@ +# Code examples + +NIR can be used to *export* or *import* models. +*Exporting* is when you convert a model from a source platform to NIR, and *importing* is when you convert a model from NIR to a target platform. +One typical workflow is to *export* a model from a simulator and *import* it to a hardware platform. + +In the menu, you see examples for how to use NIR with your favorite framework. +But note that some frameworks only support importing or exporting. + +## Writing to and reading from files with NIR +While NIR is typically integrated into your favorite framework, NIR supports writing to and reading from files directly. +This is useful when you want to send a model over email, store it for later, or share it with a colleague. + +### Writing to a file +To write a model to a file, use the `nir.write` function. +Note that this requires you to provide a NIR model, so you need to find a way to convert your model to NIR within your framework. +The `nir.write` function takes two arguments: the file path and the model to write. +```python +import nir +my_nir_graph = ... +nir.write("my_graph.nir", my_model) +``` + +### Reading from a file +To read a model from a file, use the `nir.read` function. +This function takes a single argument: the file path. +```python +import nir +imported_graph = nir.read("my_graph.nir") +``` + +This gives you a NIR model, which then needs to be converted to your framework's model format. +The NIR graph itself is just a data structure. \ No newline at end of file diff --git a/docs/source/examples/lava/test.md b/docs/source/examples/lava/test.md new file mode 100644 index 0000000..21e60f8 --- /dev/null +++ b/docs/source/examples/lava/test.md @@ -0,0 +1 @@ +# Test \ No newline at end of file diff --git a/docs/source/primitives.md b/docs/source/primitives.md index 7240268..b8d3c34 100644 --- a/docs/source/primitives.md +++ b/docs/source/primitives.md @@ -10,24 +10,25 @@ But, if you plan to execute the graph on restricted neuromorphic hardware, pleas NIR defines 16 fundamental primitives listed in the table below, which backends are free to implement as they want, leading to varying outputs across platforms. While discrepancies could be minimized by constraining implementations or making backends aware of each other's discretization choices, NIR does not do this since it is declarative, specifying only the necessary inputs and outputs. Constraining implementations would cause hardware incompatibilities and making backends aware of each other could create large O(N^2) overhead for N backends. The primitives are already computationally expressive and able to solve complex PDEs. -| Primitive | Parameters | Computation | Reset | -|-|-|-|-| -| **Input** | Input shape | - | - | -| **Output** | Output shape | - | - | -| **Affine** | $W, b$ | $ W*I + b$ | - | -| **Convolution** | $W$, Stride, Padding, Dilation, Groups, Bias | $f \star g$ | - | -| **Current-based leaky integrate-and-fire** | $\tau_\text{syn}$, $\tau_\text{mem}$, R, $v_\text{leak}$, $v_\text{thr}$, $w_\text{in}$ | **LI**_1_; **Linear**; **LIF**_2_ | $\begin{cases} v_\text{LI\_2}-v_\text{thr} & \text{Spike} \\ v & \text{else} \end{cases}$ | -| **Delay** | $\tau$ | $I(t - \tau)$ | - | -| **Flatten** | Input shape, Start dim., End dim. | - | - | -| **Integrator** | $\text{R}$ | $\dot{v} = R I$ | - | -| **Integrate-and-fire** | $\text{R}, v_\text{thr}$ | **Integrator**; **Threshold** | $\begin{cases} v-v_\text{thr} & \text{Spike} \\ v & \text{else} \end{cases}$ | -| **Leaky integrator (LI)** | $\tau, \text{R}, v_\text{leak}$ | $\tau \dot{v} = (v_\text{leak} - v) + R I$ | - | -| **Linear** | $W$ | $W I$ | - | -| **Leaky integrate-fire (LIF)** | $\tau, \text{R}, v_\text{leak}, v_\text{thr}$ | **LI**; **Threshold** | $\begin{cases} v-v_\text{thr} & \text{Spike} \\ v & \text{else} \end{cases}$ | -| **Scale** | $s$ | $s I$ | - | -| **SumPooling** | $p$ | $\sum_{j} x_j$ | | -| **AvgPooling** | $p$ | **SumPooling**; **Scale** | - | -| **Threshold** | $\theta_\text{thr}$ | $H(I - \theta_\text{thr})$ | - | +| Primitive | Parameters | Computation | Reset | +|------------------------------------|---------------------------------------------------------------------------|----------------------------------------------------------|----------------------------------------------------------------------------------------| +| **Input** | Input shape | - | - | +| **Output** | Output shape | - | - | +| **Affine** | $W, b$ | $W \cdot I + b$ | - | +| **Convolution** | $W$, Stride, Padding, Dilation, Groups, Bias | $f \star g$ | - | +| **Current-based leaky integrate-and-fire** | $\tau_\text{syn}, \tau_\text{mem}, R, v_\text{leak}, v_\text{thr}, w_\text{in}$ | **LI**; **Linear**; **LIF** | $\begin{cases} v_\text{LIF} - v_\text{thr} & \text{Spike} \\ v_\text{LIF} & \text{else} \end{cases}$ | +| **Delay** | $\tau$ | $I(t - \tau)$ | - | +| **Flatten** | Input shape, Start dim., End dim. | - | - | +| **Integrator** | $R$ | $\dot{v} = R I$ | - | +| **Integrate-and-fire** | $R, v_\text{thr}$ | **Integrator**; **Threshold** | $\begin{cases} v - v_\text{thr} & \text{Spike} \\ v & \text{else} \end{cases}$ | +| **Leaky integrator (LI)** | $\tau, R, v_\text{leak}$ | $\tau \dot{v} = (v_\text{leak} - v) + R I$ | - | +| **Linear** | $W$ | $W I$ | - | +| **Leaky integrate-fire (LIF)** | $\tau, R, v_\text{leak}, v_\text{thr}$ | **LI**; **Threshold** | $\begin{cases} v - v_\text{thr} & \text{Spike} \\ v & \text{else} \end{cases}$ | +| **Scale** | $s$ | $s I$ | - | +| **SumPooling** | $p$ | $\sum_{j} x_j$ | - | +| **AvgPooling** | $p$ | **SumPooling**; **Scale** | - | +| **Spike** | $\theta_\text{thr}$ | $\dirac(I - \theta_\text{thr})$ | - | + Each primitive is defined by their own dynamical equation, specified in the [API docs](https://nnir.readthedocs.io/en/latest/). diff --git a/docs/source/usage.md b/docs/source/usage.md index 120702f..8052cf9 100644 --- a/docs/source/usage.md +++ b/docs/source/usage.md @@ -13,6 +13,9 @@ Please refer to the **Examples** section in the sidebar for code for each suppor More code examples are available [in the repository for our paper](https://github.com/neuromorphs/NIR/tree/main/paper/). ## Example: Norse model to Sinabs Speck +This example demonstrates how to convert a Norse model to a Sinabs model and then to a Speck chip. +Note that Norse is based on PyTorch and uses [NIRTorch](#dev_pytorch) to convert PyTorch models to NIR. +You can also do this manually, by constructing your own NIR graphs as shown in our [API design documentation](#api_desige). ### Part 1: Convert Norse model to NIR ```python @@ -20,10 +23,8 @@ import torch import norse.torch as norse # Define our neural network model -model = norse.SequentialState( - norse.LIFCell(), - ... -) +model = ... + # Convert model to NIR # Note that we use some sample data to "trace" the graph in PyTorch. # You need to ensure that shape of the data fits your model @@ -44,3 +45,30 @@ dynapcnn_model = DynapcnnNetwork(sinabs_model, input_shape=sample_data.shape[-1] # Move model to chip! dynapcnn_model.to("speck2fdevkit") ``` + +## Example: Manually writing and reading NIR files +You can also manually write and read NIR files. +This is useful if you want to save a model to disk and use it later. +Or if you want to load in a model that someone else has created. + +### Writing a NIR file +[NIR consists of graphs](#primitives) that describe the structure of a neural network. +Our reference implementation uses Python to describe these graphs, so you can imagine having a graph in an object, say `nir_model`. +To write this graph to file, you can use + +```python +import nir +nir.write(nir_model, "my_model.nir") +``` + +### Reading a NIR file +Reading a NIR file is similarly easy and will give you a graph object that you can use in your code. + +```python +import nir +nir_model = nir.read("my_model.nir") +``` + +Note that the graph object (`nir_model`) doesn't do anything by itself. +You still need to convert it to a format that your hardware or simulator can understand. +Read more about this in the [Using NIR in hardware guide](#porting_nir). \ No newline at end of file