Bifrost translates models in PyTorch into SpiNNaker executables.
It is named after Bifrost, the bridge between middle earth and the realms of the gods in Norse mythology.
After training the model, when saving the StateDictionary
you need to use the utility set_parameter_buffers
included in bifrost.extract.torch.parameter_buffers
.
To load the model again, the torch.nn.Module
method load_state_dict
has to have the argument strict
set to False
.
Models are encoded by providing the Python import path to the model class as well as the shape of input tensor (here with 1 timestep, 8 batches, 2 channels, and a 640x480 image input):
bifrost model.SNNModel "(1, 8, 2, 640, 480)" > output.py
The output can now be evaluated with a specific set of parameter values from a PyTorch Lightning checkpoint file. Note that the command must run inside a SpiNNaker-friendly environment (see below):
python output.py weights.ckpt
By "SpiNNaker-friendly environment" we mean a working installation of sPyNNaker and access to a SpiNNaker machine.
See the attached Dockerfile
for a quick environment installation.
Bifrost is maintained by
- Jens E. Pedersen (@GitHub jegp), doctoral student at KTH Royal Institute of Technology, Sweden.
- Garibaldi Pineda Garcia (@GitHub chanokin), PostDoc at The University of Sussex, UK.
The project is indebted to the work by Petrut A. Bogdan on JSON conversions from ANN to SNN and to Simon Davidson for advice and supervision.
The work has received funding from the EC Horizon 2020 Framework Programme under Grant Agreements 785907 and 945539 (HBP)
LGPLv3. See LICENSE for license details.