Tutorial 2

In this tutorial, we are going to directly train a simple SNN with a single hidden layer using e-prop on the MNIST dataset, converted to a latency spike code.

Clearly, this is far from a state of the art architecture, but it still achieves 94% accuracy on MNIST.

Install

Download wheel file

[1]:
if "google.colab" in str(get_ipython()):
    !gdown 1V_GzXUDzcFz9QDIpxAD8QNEglcSipssW
    !pip install pygenn-5.0.0-cp310-cp310-linux_x86_64.whl
    %env CUDA_PATH=/usr/local/cuda

    !rm -rf /content/ml_genn
    !git clone https://github.com/genn-team/ml_genn.git --branch genn_5 -c advice.detachedHead=false
    !pip install ./ml_genn/ml_genn
Downloading...
From: https://drive.google.com/uc?id=1V_GzXUDzcFz9QDIpxAD8QNEglcSipssW
To: /content/pygenn-5.0.0-cp310-cp310-linux_x86_64.whl

0% 0.00/8.29M [00:00<?, ?B/s]

25% 2.10M/8.29M [00:00<00:00, 20.5MB/s]

100% 8.29M/8.29M [00:00<00:00, 56.7MB/s]

Processing ./pygenn-5.0.0-cp310-cp310-linux_x86_64.whl Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (1.25.2) Requirement already satisfied: deprecated in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (1.2.14) Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (5.9.5) Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated->pygenn==5.0.0) (1.14.1) pygenn is already installed with the same version as the provided wheel. Use –force-reinstall to force an installation of the wheel. env: CUDA_PATH=/usr/local/cuda Cloning into ‘ml_genn’… remote: Enumerating objects: 8243, done. remote: Counting objects: 100% (550/550), done. remote: Compressing objects: 100% (325/325), done. remote: Total 8243 (delta 274), reused 300 (delta 222), pack-reused 7693 Receiving objects: 100% (8243/8243), 37.42 MiB | 19.23 MiB/s, done. Resolving deltas: 100% (5506/5506), done. Processing ./ml_genn/ml_genn

Preparing metadata (setup.py) … done

Requirement already satisfied: pygenn<6.0.0,>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (5.0.0) Collecting enum-compat (from ml-genn==2.1.0)

Using cached enum_compat-0.0.3-py3-none-any.whl (1.3 kB)

Requirement already satisfied: tqdm>=4.27.0 in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (4.66.2) Requirement already satisfied: deprecated in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (1.2.14) Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from pygenn<6.0.0,>=5.0.0->ml-genn==2.1.0) (1.25.2) Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from pygenn<6.0.0,>=5.0.0->ml-genn==2.1.0) (5.9.5) Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated->ml-genn==2.1.0) (1.14.1) Building wheels for collected packages: ml-genn

Building wheel for ml-genn (setup.py) … done Created wheel for ml-genn: filename=ml_genn-2.1.0-py3-none-any.whl size=109412 sha256=7f4edda1b1644da8f8847b7168d4cd367a68435740436bed18cbdd81b77b258c Stored in directory: /tmp/pip-ephem-wheel-cache-q0xomaml/wheels/3f/cf/27/0e9dec4bb1be2afac4b38c2dfb4ce0bc164ce1ecb32b6f91b8

Successfully built ml-genn Installing collected packages: enum-compat, ml-genn Successfully installed enum-compat-0.0.3 ml-genn-2.1.0 </pre>

0% 0.00/8.29M [00:00<?, ?B/s]

25% 2.10M/8.29M [00:00<00:00, 20.5MB/s]

100% 8.29M/8.29M [00:00<00:00, 56.7MB/s]

Processing ./pygenn-5.0.0-cp310-cp310-linux_x86_64.whl Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (1.25.2) Requirement already satisfied: deprecated in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (1.2.14) Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (5.9.5) Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated->pygenn==5.0.0) (1.14.1) pygenn is already installed with the same version as the provided wheel. Use –force-reinstall to force an installation of the wheel. env: CUDA_PATH=/usr/local/cuda Cloning into ‘ml_genn’{ldots} remote: Enumerating objects: 8243, done. remote: Counting objects: 100% (550/550), done. remote: Compressing objects: 100% (325/325), done. remote: Total 8243 (delta 274), reused 300 (delta 222), pack-reused 7693 Receiving objects: 100% (8243/8243), 37.42 MiB | 19.23 MiB/s, done. Resolving deltas: 100% (5506/5506), done. Processing ./ml_genn/ml_genn

Preparing metadata (setup.py) {ldots} done

Requirement already satisfied: pygenn<6.0.0,>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (5.0.0) Collecting enum-compat (from ml-genn==2.1.0)

Using cached enum_compat-0.0.3-py3-none-any.whl (1.3 kB)

Requirement already satisfied: tqdm>=4.27.0 in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (4.66.2) Requirement already satisfied: deprecated in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (1.2.14) Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from pygenn<6.0.0,>=5.0.0->ml-genn==2.1.0) (1.25.2) Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from pygenn<6.0.0,>=5.0.0->ml-genn==2.1.0) (5.9.5) Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated->ml-genn==2.1.0) (1.14.1) Building wheels for collected packages: ml-genn

Building wheel for ml-genn (setup.py) {ldots} done Created wheel for ml-genn: filename=ml_genn-2.1.0-py3-none-any.whl size=109412 sha256=7f4edda1b1644da8f8847b7168d4cd367a68435740436bed18cbdd81b77b258c Stored in directory: /tmp/pip-ephem-wheel-cache-q0xomaml/wheels/3f/cf/27/0e9dec4bb1be2afac4b38c2dfb4ce0bc164ce1ecb32b6f91b8

Successfully built ml-genn Installing collected packages: enum-compat, ml-genn Successfully installed enum-compat-0.0.3 ml-genn-2.1.0 end{sphinxVerbatim}

0% 0.00/8.29M [00:00<?, ?B/s]

25% 2.10M/8.29M [00:00<00:00, 20.5MB/s]

100% 8.29M/8.29M [00:00<00:00, 56.7MB/s]

Processing ./pygenn-5.0.0-cp310-cp310-linux_x86_64.whl Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (1.25.2) Requirement already satisfied: deprecated in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (1.2.14) Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from pygenn==5.0.0) (5.9.5) Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated->pygenn==5.0.0) (1.14.1) pygenn is already installed with the same version as the provided wheel. Use –force-reinstall to force an installation of the wheel. env: CUDA_PATH=/usr/local/cuda Cloning into ‘ml_genn’… remote: Enumerating objects: 8243, done. remote: Counting objects: 100% (550/550), done. remote: Compressing objects: 100% (325/325), done. remote: Total 8243 (delta 274), reused 300 (delta 222), pack-reused 7693 Receiving objects: 100% (8243/8243), 37.42 MiB | 19.23 MiB/s, done. Resolving deltas: 100% (5506/5506), done. Processing ./ml_genn/ml_genn

Preparing metadata (setup.py) … [?25l[?25hdone

Requirement already satisfied: pygenn<6.0.0,>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (5.0.0) Collecting enum-compat (from ml-genn==2.1.0)

Using cached enum_compat-0.0.3-py3-none-any.whl (1.3 kB)

Requirement already satisfied: tqdm>=4.27.0 in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (4.66.2) Requirement already satisfied: deprecated in /usr/local/lib/python3.10/dist-packages (from ml-genn==2.1.0) (1.2.14) Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from pygenn<6.0.0,>=5.0.0->ml-genn==2.1.0) (1.25.2) Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from pygenn<6.0.0,>=5.0.0->ml-genn==2.1.0) (5.9.5) Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated->ml-genn==2.1.0) (1.14.1) Building wheels for collected packages: ml-genn

Building wheel for ml-genn (setup.py) … [?25l[?25hdone Created wheel for ml-genn: filename=ml_genn-2.1.0-py3-none-any.whl size=109412 sha256=7f4edda1b1644da8f8847b7168d4cd367a68435740436bed18cbdd81b77b258c Stored in directory: /tmp/pip-ephem-wheel-cache-q0xomaml/wheels/3f/cf/27/0e9dec4bb1be2afac4b38c2dfb4ce0bc164ce1ecb32b6f91b8

Successfully built ml-genn Installing collected packages: enum-compat, ml-genn Successfully installed enum-compat-0.0.3 ml-genn-2.1.0

Install MNIST package

[2]:
!pip install mnist
Requirement already satisfied: mnist in /usr/local/lib/python3.10/dist-packages (0.2.2)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from mnist) (1.25.2)

Build model

Import standard modules and required mlGeNN classes

[3]:
import mnist
import numpy as np
import matplotlib.pyplot as plt

from ml_genn import InputLayer, Layer, SequentialNetwork
from ml_genn.callbacks import Checkpoint
from ml_genn.compilers import EPropCompiler, InferenceCompiler
from ml_genn.connectivity import Dense,FixedProbability
from ml_genn.initializers import Normal
from ml_genn.neurons import LeakyIntegrate, LeakyIntegrateFire, SpikeInput
from ml_genn.serialisers import Numpy

from ml_genn.utils.data import (calc_latest_spike_time, log_latency_encode_data)

from ml_genn.compilers.eprop_compiler import default_params

##Parameters

Define some model parameters

[4]:
NUM_INPUT = 28 * 28
NUM_HIDDEN = 128
NUM_OUTPUT = 10
BATCH_SIZE = 128

Latency encoding

There are numerous ways to encode images using spikes but here we are going to emit a single spike for each neuron at a time calculated as follows from the pixel grayscale \(x\): :nbsphinx-math:`begin{align}
T(x) = begin{cases}

tau_text{eff} logleft(frac{x}{x-theta} right) & x > theta\ infty & otherwise\

end{cases}

end{align}` where \(\tau_\text{eff}=20\text{ms}\) and \(\theta=51\).

[5]:
train_spikes = log_latency_encode_data(mnist.train_images(), 20.0, 51)

Network definition

Because our network is entirely feedforward, we can define it as a SequentialNetwork where each layer is automatically connected to the previous layer. As we have converted the MNIST dataset to spikes, we will use a SpikeInput to inject these directly into the network. For our hidden layer we are going to use standard Leaky integrate-and-fire neurons as this task does not require more computationally expensive adaptive LIF neurons. Finally, we are going to use a non-spiking output layer and read classifications out of this by determining the maximum of the summed membrane voltages of the output neurons.

[6]:
# Create sequential model
serialiser = Numpy("latency_mnist_checkpoints")
network = SequentialNetwork(default_params)
with network:
    # Populations
    input = InputLayer(SpikeInput(max_spikes=BATCH_SIZE * NUM_INPUT),
                                  NUM_INPUT)
    hidden = Layer(Dense(Normal(sd=1.0 / np.sqrt(NUM_INPUT))),
                   LeakyIntegrateFire(v_thresh=0.61, tau_mem=20.0,
                                      tau_refrac=5.0),
                   NUM_HIDDEN)
    output = Layer(Dense(Normal(sd=1.0 / np.sqrt(NUM_HIDDEN))),
                   LeakyIntegrate(tau_mem=20.0, readout="sum_var"),
                   NUM_OUTPUT)

Compilation

In mlGeNN, in order to turn an abstract network description into something that can actually be used for training or inference you use a compiler class. Here, we use the EPropCompiler to train with e-prop and specify batch size and how many timesteps to evaluate each example for as well as choosing our optimiser and loss function. Because this is a classification task, we want to use cross-entropy loss and, because our labels are specified in this way (rather than e.g. one-hot encoded), we use the sparse catgorical variant.

[7]:
max_example_timesteps = int(np.ceil(calc_latest_spike_time(train_spikes)))
compiler = EPropCompiler(example_timesteps=max_example_timesteps,
                         losses="sparse_categorical_crossentropy",
                         optimiser="adam", batch_size=BATCH_SIZE)
compiled_net = compiler.compile(network)

Training

Now we will train the model for 10 epochs using our compiled network. To verify it’s performance we take 10% of the training data as a validation split and add an additional callback to checkpoint weights every epoch.

[8]:
with compiled_net:
    # Evaluate model on numpy dataset
    callbacks = ["batch_progress_bar", Checkpoint(serialiser)]
    compiled_net.train({input: train_spikes},
                       {output: mnist.train_labels()},
                       num_epochs=10, shuffle=True,
                       validation_split=0.1,
                       callbacks=callbacks)

Evaluate

Load weights checkpointed from last epoch:

[9]:
network.load((9,), serialiser)

Create an InferenceCompiler and compile network for inference:

[10]:
compiler = InferenceCompiler(evaluate_timesteps=max_example_timesteps,
                             batch_size=BATCH_SIZE)
compiled_net = compiler.compile(network)

Encode test set using the same log-latency encoding and evaluate it:

[11]:
test_spikes = log_latency_encode_data(mnist.test_images(), 20.0, 51)
with compiled_net:
    compiled_net.evaluate({input: test_spikes},
                          {output: mnist.test_labels()})