Overview
NIRTorch is a powerful extension for PyTorch that bridges the gap between PyTorch-based Spiking Neural Network (SNN) frameworks and the Neuromorphic Intermediate Representation (NIR). Its primary purpose is to provide a simple and robust way to convert torch.nn.Module objects into NIR graphs, and vice-versa.
By leveraging torch.fx for symbolic tracing, NIRTorch can inspect a PyTorch model’s computational graph and map its components to the corresponding NIR primitives. This allows developers to design and train SNNs within their favorite PyTorch-based environment (like snnTorch or Norse) and then export them to a standardized format that can be deployed across various neuromorphic hardware platforms and simulators that support NIR.
Key features include:
- A simple
to_nir()function to convert atorch.nn.Moduleto a NIR graph. - A
from_nir()function to load a NIR graph back into a runnable PyTorch model. - A framework-agnostic design that allows any PyTorch-based SNN library to integrate with NIR by providing a simple mapping of its modules.
