jeshraghian / snntorch

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

Update NIR <> snnTorch #246

Closed stevenabreu7 closed 8 months ago

stevenabreu7 commented 11 months ago

This PR will:

For RNNs to work, make sure to use the right version of NIRTorch (current PR: https://github.com/neuromorphs/NIRTorch/pull/12)

stevenabreu7 commented 8 months ago

after some (way too much) time, this PR is ready for review! @jeshraghian @ahenkes1 can you advise on next steps? add tests? docs? how can I help? cheers!

jeshraghian commented 8 months ago

This is wonderful, thanks @stevenabreu7 !

docs

In terms of docs, we should generate docstrings that describe the arguments of those two import/export functions with a minimalist example.

Any caveats can be added as bullet points in the docstrings, e.g., types of neurons supported, layers, neuron arguments supported.

I can create a skeleton, though would you be able to fill the meat?

Once we have that, I can also put together some ipynb stuff based on these docstrings.

tests It'd be worth adding a test for our leaky neurons when the reset_delay is active. I can handle this.

I'm not too sure how we can do unit tests for the actual nir-dependent stuff, though... any ideas? Could it be a matter of having a minimalist set of models, passing them in, making sure their outputs have the arguments (just an assertion the arguments exist. Checking its values is overkill and left to nir). Thoughts?

stevenabreu7 commented 8 months ago

sounds good! I added docstrings, lmk if they're okay like that. Also added suggestions for how to test in test_nir.py, what do you think? if you're okay with it, I'll fill in the rest @jeshraghian

jeshraghian commented 8 months ago

Alright, looking great @stevenabreu7!

Some of the syntax from the leaky neuron was modified in another PR to make it compatible with torch.compile() so there were a few conflicts + failing tests that I've now fixed and pushed to your branch.

Your suggestions for tests look great, as do the docstrings. Once the tests are done then it's good to merge.

stevenabreu7 commented 8 months ago

thank you @jeshraghian! added test and some notes on the rleaky that's not currently supported. I can make a new PR for this, wouldn't take long, but can we merge already so that it runs smoothly for the NIR demo tomorrow? (otherwise it's a bit of a pain to use this branch of snntorch for the demo)

stevenabreu7 commented 8 months ago

added an issue for RLeaky support - will update the NIR support in a fresh PR in the near future

jeshraghian commented 8 months ago

Just merged! The demo will still need to be from the latest version that is installed from the source rather than pip installing. Though I can do a version increment in the morning, if that helps. Lmk.