genn-team / ml_genn

A library for deep learning with Spiking Neural Networks (SNN).
https://ml-genn.readthedocs.io
GNU Lesser General Public License v2.1
24 stars 7 forks source link

Data parallel training support #79

Closed neworderofjamie closed 11 months ago

neworderofjamie commented 1 year ago

Pleasingly, this was actually very easy to do! Basically:

Other than that, it's all just passing the communicator around and a few places where the 'full' batch size is used rather than the scaled down one e.g. in the EventProp compiler to scale stuff. I've also added a couple of additional examples (at some point I need to tidy the examples up a bit) which demonstrate how you need to change your code to run across multiple GPUs - mostly just splitting the dataset and turning off progress bars etc apart from on the first rank.