jeshraghian / snntorch

Deep and online learning with spiking neural networks in Python
https://snntorch.readthedocs.io/en/latest/
MIT License
1.17k stars 202 forks source link

Missing Conv2D and AveragePooling2D modules in export/import to NIR #304

Open SirineArfa opened 3 months ago

SirineArfa commented 3 months ago

Description

The currently supported modules in the snntorch's _exportnir.py are [Leaky, Linear, Synaptic, RLeaky, RSynaptic]. In my case I need to convert an SNN model with Conv2D, Leaky, Linear and AveragePooling2D layers to NIR.

What could be done

Two possible options to solve this:

  1. Naively add these two torch.nn layer operations to the _snntorch/exportnir.py function
  2. Centralizing NIRTorch functions _fromnir.py and _tonir.py to support the NIR conversion of torch.nn layers regardless of the used framework (snntorch,norse,sinabs...)