Copyright (C) 2021 Intel Corporation SPDX-License-Identifier: BSD-3-Clause See:

MNIST Digit Classification with Lava

Motivation: In this tutorial, we will build a Lava Process for an MNIST classifier, using the Lava Processes for LIF neurons and Dense connectivity. The tutorial is useful to get started with Lava in a few minutes.

This tutorial assumes that you:

This tutorial gives a bird’s-eye view of

Our MNIST Classifier

In this tutorial, we will build a multi-layer feed-forward classifier without any convolutional layers. The architecture is shown below.

Important Note:

The classifier is a simple feed-forward model using pre-trained network parameters (weights and biases). It illustrates how to build, compile and run a functional model in Lava. Please refer to Lava-DL to understand how to train deep networks with Lava.

MNIST Process AA

The 3 Processes shown above are: - SpikeInput Process - generates spikes via integrate and fire dynamics, using the image input - ImageClassifier Process - encapsulates feed-forward architecture of Dense connectivity and LIF neurons - Output Process - accumulates output spikes from the feed-forward process and infers the class label

General Imports

import os
import numpy as np
import typing as ty

Lava Processes

Below we create the Lava Process classes. We need to define only the structure of the process here. The details about how the Process will be executed are specified in the ProcessModels below.

As mentioned above, we define Processes for - converting input images to binary spikes from those biases (SpikeInput), - the 3-layer fully connected feed-forward network (MnistClassifier) - accumulating the output spikes and inferring the class for an input image (OutputProcess)

# Import Process level primitives
from lava.magma.core.process.process import AbstractProcess
from lava.magma.core.process.variable import Var
from lava.magma.core.process.ports.ports import InPort, OutPort
class SpikeInput(AbstractProcess):
    """Reads image data from the MNIST dataset and converts it to spikes.
    The resulting spike rate is proportional to the pixel value."""

    def __init__(self,
                 vth: int,
                 num_images: ty.Optional[int] = 25,
                 num_steps_per_image: ty.Optional[int] = 128):
        shape = (784,)
        self.spikes_out = OutPort(shape=shape)  # Input spikes to the classifier
        self.label_out = OutPort(shape=(1,))  # Ground truth labels to OutputProc
        self.num_images = Var(shape=(1,), init=num_images)
        self.num_steps_per_image = Var(shape=(1,), init=num_steps_per_image)
        self.input_img = Var(shape=shape)
        self.ground_truth_label = Var(shape=(1,))
        self.v = Var(shape=shape, init=0)
        self.vth = Var(shape=(1,), init=vth)

class ImageClassifier(AbstractProcess):
    """A 3 layer feed-forward network with LIF and Dense Processes."""

    def __init__(self, trained_weights_path: str):

        # Using pre-trained weights and biases
        real_path_trained_wgts = os.path.realpath(trained_weights_path)

        wb_list = np.load(real_path_trained_wgts, encoding='latin1', allow_pickle=True)
        w0 = wb_list[0].transpose().astype(np.int32)
        w1 = wb_list[2].transpose().astype(np.int32)
        w2 = wb_list[4].transpose().astype(np.int32)
        b1 = wb_list[1].astype(np.int32)
        b2 = wb_list[3].astype(np.int32)
        b3 = wb_list[5].astype(np.int32)

        self.spikes_in = InPort(shape=(w0.shape[1],))
        self.spikes_out = OutPort(shape=(w2.shape[0],))
        self.w_dense0 = Var(shape=w0.shape, init=w0)
        self.b_lif1 = Var(shape=(w0.shape[0],), init=b1)
        self.w_dense1 = Var(shape=w1.shape, init=w1)
        self.b_lif2 = Var(shape=(w1.shape[0],), init=b2)
        self.w_dense2 = Var(shape=w2.shape, init=w2)
        self.b_output_lif = Var(shape=(w2.shape[0],), init=b3)

        # Up-level currents and voltages of LIF Processes
        # for resetting (see at the end of the tutorial)
        self.lif1_u = Var(shape=(w0.shape[0],), init=0)
        self.lif1_v = Var(shape=(w0.shape[0],), init=0)
        self.lif2_u = Var(shape=(w1.shape[0],), init=0)
        self.lif2_v = Var(shape=(w1.shape[0],), init=0)
        self.oplif_u = Var(shape=(w2.shape[0],), init=0)
        self.oplif_v = Var(shape=(w2.shape[0],), init=0)

class OutputProcess(AbstractProcess):
    """Process to gather spikes from 10 output LIF neurons and interpret the
    highest spiking rate as the classifier output"""

    def __init__(self, **kwargs):
        shape = (10,)
        n_img = kwargs.pop('num_images', 25)
        self.num_images = Var(shape=(1,), init=n_img)
        self.spikes_in = InPort(shape=shape)
        self.label_in = InPort(shape=(1,))
        self.spikes_accum = Var(shape=shape)  # Accumulated spikes for classification
        self.num_steps_per_image = Var(shape=(1,), init=128)
        self.pred_labels = Var(shape=(n_img,))
        self.gt_labels = Var(shape=(n_img,))

ProcessModels for Python execution

# Import parent classes for ProcessModels
from lava.magma.core.model.sub.model import AbstractSubProcessModel
from import PyLoihiProcessModel

# Import ProcessModel ports, data-types
from import PyInPort, PyOutPort
from import LavaPyType

# Import execution protocol and hardware resources
from lava.magma.core.sync.protocols.loihi_protocol import LoihiProtocol
from lava.magma.core.resources import CPU

# Import decorators
from lava.magma.core.decorator import implements, requires

# Import MNIST dataset
from lava.utils.dataloader.mnist import MnistDataset

Decorators for ProcessModels: - @implements: associates a ProcessModel with a Process through the argument proc. Using protocol argument, we will specify the synchronization protocol used by the ProcessModel. In this tutorial, all ProcessModels execute according to the LoihiProtocol. Which means, similar to the Loihi chip, each time-step is divided into spiking, pre-management, post-management, and learning phases. It is necessary to specify behaviour of a ProcessModel during the spiking phase using run_spk function. Other phases are optional. - @requires: specifies the hardware resource on which a ProcessModel will be executed. In this tutorial, we will execute all ProcessModels on a CPU.

SpikingInput ProcessModel

@implements(proc=SpikeInput, protocol=LoihiProtocol)
class PySpikeInputModel(PyLoihiProcessModel):
    num_images: int = LavaPyType(int, int, precision=32)
    spikes_out: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, bool, precision=1)
    label_out: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, np.int32,
    num_steps_per_image: int = LavaPyType(int, int, precision=32)
    input_img: np.ndarray = LavaPyType(np.ndarray, int, precision=32)
    ground_truth_label: int = LavaPyType(int, int, precision=32)
    v: np.ndarray = LavaPyType(np.ndarray, int, precision=32)
    vth: int = LavaPyType(int, int, precision=32)

    def __init__(self, proc_params):
        self.mnist_dataset = MnistDataset()
        self.curr_img_id = 0

    def post_guard(self):
        """Guard function for PostManagement phase.
        if self.time_step % self.num_steps_per_image == 1:
            return True
        return False

    def run_post_mgmt(self):
        """Post-Management phase: executed only when guard function above
        returns True.
        img = self.mnist_dataset.images[self.curr_img_id]
        self.ground_truth_label = self.mnist_dataset.labels[self.curr_img_id]
        self.input_img = img.astype(np.int32) - 127
        self.v = np.zeros(self.v.shape)
        self.curr_img_id += 1

    def run_spk(self):
        """Spiking phase: executed unconditionally at every time-step
        self.v[:] = self.v + self.input_img
        s_out = self.v > self.vth
        self.v[s_out] = 0  # reset voltage to 0 after a spike

ImageClassifier ProcessModel

Notice that the following process model is further decomposed into sub-Processes, which implement LIF neural dynamics and Dense connectivity. We will not go into the details of how these are implemented in this tutorial.

Also notice that a SubProcessModel does not actually contain any concrete execution. This is handled by the ProcessModels of the constituent Processes.

from lava.proc.lif.process import LIF
from lava.proc.dense.process import Dense

class PyImageClassifierModel(AbstractSubProcessModel):
    def __init__(self, proc):
        self.dense0 = Dense(weights=proc.w_dense0.init)
        self.lif1 = LIF(shape=(64,), bias_mant=proc.b_lif1.init, vth=400,
                        dv=0, du=4095)
        self.dense1 = Dense(weights=proc.w_dense1.init)
        self.lif2 = LIF(shape=(64,), bias_mant=proc.b_lif2.init, vth=350,
                        dv=0, du=4095)
        self.dense2 = Dense(weights=proc.w_dense2.init)
        self.output_lif = LIF(shape=(10,), bias_mant=proc.b_output_lif.init,
                              vth=1, dv=0, du=4095)


        # Create aliases of SubProcess variables

OutputProcess ProcessModel

@implements(proc=OutputProcess, protocol=LoihiProtocol)
class PyOutputProcessModel(PyLoihiProcessModel):
    label_in: PyInPort = LavaPyType(PyInPort.VEC_DENSE, int, precision=32)
    spikes_in: PyInPort = LavaPyType(PyInPort.VEC_DENSE, bool, precision=1)
    num_images: int = LavaPyType(int, int, precision=32)
    spikes_accum: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=32)
    num_steps_per_image: int = LavaPyType(int, int, precision=32)
    pred_labels: np.ndarray = LavaPyType(np.ndarray, int, precision=32)
    gt_labels: np.ndarray = LavaPyType(np.ndarray, int, precision=32)

    def __init__(self, proc_params):
        self.current_img_id = 0

    def post_guard(self):
        """Guard function for PostManagement phase.
        if self.time_step % self.num_steps_per_image == 0 and \
                self.time_step > 1:
            return True
        return False

    def run_post_mgmt(self):
        """Post-Management phase: executed only when guard function above
        returns True.
        gt_label = self.label_in.recv()
        pred_label = np.argmax(self.spikes_accum)
        self.gt_labels[self.current_img_id] = gt_label
        self.pred_labels[self.current_img_id] = pred_label
        self.current_img_id += 1
        self.spikes_accum = np.zeros_like(self.spikes_accum)

    def run_spk(self):
        """Spiking phase: executed unconditionally at every time-step
        spk_in = self.spikes_in.recv()
        self.spikes_accum = self.spikes_accum + spk_in

Connecting Processes

num_images = 25
num_steps_per_image = 128

# Create Process instances
spike_input = SpikeInput(vth=1,
mnist_clf = ImageClassifier(
    trained_weights_path=os.path.join('.', 'mnist_pretrained.npy'))
output_proc = OutputProcess(num_images=num_images)

# Connect Processes
# Connect Input directly to Output for ground truth labels

If you receive an UnpicklingError when instantiating the ImageClassifier, make sure to download the pretrained weights from GitHub LFS in the current directory using:

git lfs fetch
git lfs pull

Execution and results

Below, we will run the classifier process in a loop of num_images number of iterations. Each iteration will run the Process for num_steps_per_image number of time-steps.

We take this approach to clear the neural states of all three LIF layers inside the classifier after every image. We need to clear the neural states, because the network parameters were trained assuming clear neural states for each inference.

Note: Below we have used Var.set() function to set the values of internal state variables. The same behaviour can be achieved by using RefPorts. See the RefPorts tutorial to learn more about how to use RefPorts to access internal state variables of Lava Processes.

from lava.magma.core.run_conditions import RunSteps
from lava.magma.core.run_configs import Loihi1SimCfg

# Loop over all images
for img_id in range(num_images):
    print(f"\rCurrent image: {img_id+1}", end="")

    # Run each image-inference for fixed number of steps

    # Reset internal neural state of LIF neurons
    mnist_clf.lif1_u.set(np.zeros((64,), dtype=np.int32))
    mnist_clf.lif1_v.set(np.zeros((64,), dtype=np.int32))
    mnist_clf.lif2_u.set(np.zeros((64,), dtype=np.int32))
    mnist_clf.lif2_v.set(np.zeros((64,), dtype=np.int32))
    mnist_clf.oplif_u.set(np.zeros((10,), dtype=np.int32))
    mnist_clf.oplif_v.set(np.zeros((10,), dtype=np.int32))

# Gather ground truth and predictions before stopping exec
ground_truth = output_proc.gt_labels.get().astype(np.int32)
predictions = output_proc.pred_labels.get().astype(np.int32)

# Stop the execution

accuracy = np.sum(ground_truth==predictions)/ground_truth.size * 100

print(f"\nGround truth: {ground_truth}\n"
      f"Predictions : {predictions}\n"
      f"Accuracy    : {accuracy}")
Current image: 25
Ground truth: [5 0 4 1 9 2 1 3 1 4 3 5 3 6 1 7 2 8 6 9 4 0 9 1 1]
Predictions : [3 0 4 1 4 2 1 3 1 4 3 5 3 6 1 7 2 8 5 9 4 0 9 1 1]
Accuracy    : 88.0

Important Note:

The classifier is a simple feed-forward model using pre-trained network parameters (weights and biases). It illustrates how to build, compile and run a functional model in Lava. Please refer to Lava-DL to understand how to train deep networks with Lava.

How to learn more?