Convert to Lava for Loihi Deployment

Author: fangwei123456

Thanks to AllenYolk and banzhuangonglxh for their contributions to lava_exchange

Introduction of Lava

Lava is a neuromorphic computing framework, which is mainly developed by Intel and supports deploying on Intel Loihi. Lava provides a sub-package Lava DL for deep learning, which can be used to build and train deep SNNs.

To deploy SNNs on Loihi, we need to use Lava. SpikingJelly provides conversion modules to convert the SNN trained by SpikingJelly to the Lava SNN format. And then we can run this SNN on Loihi. The workflow is:

SpikingJelly -> Lava DL -> Lava -> Loihi

The modules related to Lava are defined in spikingjelly.activation_based.lava_exchange.

Basic Conversion

Data Format Conversion

The default data format in Lava DL is shape = [N, *, T], where N is the batch dimension and T is the time-step dimension. However, the module of SpikingJelly in multi-step mode (step_mode = 'm') uses the data format as shape = [T, N, *]. Thus, lava_exchange provides two conversion functions, TNX_to_NXT and NXT_to_TNX for conversion between two formats. Here is an example:

import torch
from spikingjelly.activation_based import lava_exchange

T = 6
N = 4
C = 2

x_seq = torch.rand([T, N, C])

x_seq_la = lava_exchange.TNX_to_NXT(x_seq)
print(f'x_seq_la.shape=[N, C, T]={x_seq_la.shape}')

x_seq_sj = lava_exchange.NXT_to_TNX(x_seq_la)
print(f'x_seq_sj.shape=[T, N, C]={x_seq_sj.shape}')

The outputs are:

x_seq_la.shape=[N, C, T]=torch.Size([4, 2, 6])
x_seq_sj.shape=[T, N, C]=torch.Size([6, 4, 2])

Neuron Conversion

Neurons in SpikingJelly can be converted to neurons in Lava DL. Due to the limited time and energy of developers, SpikingJelly only supports the IF neuron and the LIF neuron, which are two of the most popular neurons in spiking deep learning. Other neurons will be considered to add according to user requirements.

We can use to_lava_neuron to convert. Here is an example:

import torch
from spikingjelly.activation_based import lava_exchange, neuron

if_sj = neuron.IFNode(v_threshold=1., v_reset=0., step_mode='m')
if_la = lava_exchange.to_lava_neuron(if_sj)

T = 8
N = 2
C = 1

x_seq_sj = torch.rand([T, N, C])
x_seq_la = lava_exchange.TNX_to_NXT(x_seq_sj)

print('output of sj(reshaped to NXT):\n', lava_exchange.TNX_to_NXT(if_sj(x_seq_sj)))
print('output of lava:\n', if_la(x_seq_la))

The outputs are:

output of sj(reshaped to NXT):
tensor([[[0., 0., 1., 0., 1., 0., 0., 0.]],

        [[0., 1., 0., 1., 0., 1., 0., 1.]]])
output of lava:
tensor([[[0., 0., 1., 0., 1., 0., 0., 0.]],

        [[0., 1., 0., 1., 0., 1., 0., 1.]]])

Here is an example of using the LIF neuron:

import torch
from spikingjelly.activation_based import lava_exchange, neuron

if_sj = neuron.LIFNode(tau=50., decay_input=False, v_threshold=1., v_reset=0., step_mode='m')
if_la = lava_exchange.to_lava_neuron(if_sj)

T = 8
N = 2
C = 1

x_seq_sj = torch.rand([T, N, C])
x_seq_la = lava_exchange.TNX_to_NXT(x_seq_sj)

print('output of sj:\n', lava_exchange.TNX_to_NXT(if_sj(x_seq_sj)))
print('output of lava:\n', if_la(x_seq_la))

The outputs are:

output of sj:
tensor([[[0., 1., 0., 1., 0., 0., 1., 0.]],

        [[0., 0., 1., 0., 0., 1., 0., 1.]]])
output of lava:
tensor([[[0., 1., 0., 1., 0., 0., 1., 0.]],

        [[0., 0., 1., 0., 0., 1., 0., 1.]]])

Synapse Conversion

The frequently-used convolutional layer, linear layer, and pooling layer can be converted. Note that

  • bias is not supported

  • Lava only supports sum pooling, which can be regarded as average pooling without average

Here is an example:

from spikingjelly.activation_based import lava_exchange, layer

conv = layer.Conv2d(3, 4, kernel_size=3, stride=1, bias=False)
fc = layer.Linear(4, 2, bias=False)
ap = layer.AvgPool2d(2, 2)

conv_la = lava_exchange.conv2d_to_lava_synapse_conv(conv)
fc_la = lava_exchange.linear_to_lava_synapse_dense(fc)
sp_la = lava_exchange.avgpool2d_to_lava_synapse_pool(ap)

print(f'conv_la={conv_la}')
print(f'fc_la={fc_la}')
print(f'sp_la={sp_la}')

The outputs are:

WARNING:root:The lava slayer pool layer applies sum pooling, rather than average pooling. `avgpool2d_to_lava_synapse_pool` will return a sum pooling layer.
conv_la=Conv(3, 4, kernel_size=(3, 3, 1), stride=(1, 1, 1), bias=False)
fc_la=Dense(4, 2, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
sp_la=Pool(1, 1, kernel_size=(2, 2, 1), stride=(2, 2, 1), bias=False)

Almost all synapses in Lava DL are based on torch.nn.Conv3d. Thus, when we print them, we will find that kernel_size and stride are tuples with three elements.

BlockContainer

The workflow for using Lava DL is:

  1. using Blocks in Lava DL to build and train the deep SNN

  2. exporting the SNN to the hdf5 file

  3. using Lava to read the hdf5 file and rebuild the SNN, then the SNN can run on Loihi or the CPU-simulated Loihi

For more details, please refer to Lava: Deep Learning.

Blocks can be regarded as the ensemble of a synapse layer and a neuron layer. For example, lava.lib.dl.slayer.block.cuba.Conv is composed of a convolutional layer and a CUDA LIF neuron layer.

Note that Blocks is designed for SNN deployment. Thus, synapses and neuronal dynamics are quantized in Blocks. Thus, Blocks is not a simple synapse + neuron ``, but ``quantize(synapse) + quantize(neuron).

SpikingJelly provides BlockContainer to mimic Blocks in Lava. The features of BlockContainer are as follows:

  • supports for surrogate gradient training

  • synapses and neuronal dynamics are quantized

  • the outputs are identical to Blocks of Lava DL when giving the same inputs

  • supports for converting to lava.lib.dl.slayer.block

For the moment, BlockContainer only supports for lava_exchange.CubaLIFNode. But it also supports for converting IFNode or LIFNode in init args to CubaLIFNode. Here is an example:

from spikingjelly.activation_based import lava_exchange, layer, neuron

fc_block_sj = lava_exchange.BlockContainer(
    synapse=layer.Linear(8, 1, bias=False),
    neu=neuron.IFNode(),
    step_mode='m'
)

print('fc_block_sj=\n', fc_block_sj)

fc_block_la = fc_block_sj.to_lava_block()
print('fc_block_la=\n', fc_block_la)

The outputs are:

fc_block_sj=
BlockContainer(
(synapse): Linear(in_features=8, out_features=1, bias=False)
(neuron): CubaLIFNode(
    v_threshold=1.0, v_reset=0.0, detach_reset=False, step_mode=m, backend=torch
    (surrogate_function): Sigmoid(alpha=4.0, spiking=True)
)
)
fc_block_la=
Dense(
(neuron): Neuron()
(synapse): Dense(8, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
)

MNIST CSNN Example

Now let us train a spiking convolutional SNN for classifying MNIST, and then convert this network to Lava DL format.

The SNN is defined as:

class MNISTNet(nn.Module):
    def __init__(self, channels: int = 16):
        super().__init__()
        self.conv_fc = nn.Sequential(
            lava_exchange.BlockContainer(
                nn.Conv2d(1, channels, kernel_size=3, stride=1, padding=1, bias=False),
                neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)
            ),

            lava_exchange.BlockContainer(
                nn.Conv2d(channels, channels, kernel_size=2, stride=2, bias=False),
                neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)
            ),
            # 14 * 14

            lava_exchange.BlockContainer(
                nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1, bias=False),
                neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)
            ),

            lava_exchange.BlockContainer(
                nn.Conv2d(channels, channels, kernel_size=2, stride=2, bias=False),
                neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)
            ),

            # 7 * 7

            lava_exchange.BlockContainer(
                nn.Flatten(),
                None
            ),
            lava_exchange.BlockContainer(
                nn.Linear(channels * 7 * 7, 128, bias=False),
                neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)
            ),

            lava_exchange.BlockContainer(
                nn.Linear(128, 10, bias=False),
                neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)
            ),
        )

    def forward(self, x):
        return self.conv_fc(x)

We add a conversion function to convert the SNN to Lava DL format, which can be used after training:

def to_lava(self):
    ret = []

    for i in range(self.conv_fc.__len__()):
        m = self.conv_fc[i]
        if isinstance(m, lava_exchange.BlockContainer):
            ret.append(m.to_lava_block())

    return nn.Sequential(*ret)

Then, we train this SNN. The training process has no much difference from other SNNs. Note that the quantization inside lava_exchange.BlockContainer will reduce accuracy. An example of the training codes is:

encoder = encoding.PoissonEncoder(step_mode='m')
# ...
for img, label in train_data_loader:
    optimizer.zero_grad()
    img = img.to(args.device)
    label = label.to(args.device)
    img = img.unsqueeze(0).repeat(args.T, 1, 1, 1, 1)

    fr = net(encoder(img)).mean(0)
    loss = F.cross_entropy(fr, label)
    loss.backward()
    optimizer.step()
    # ...

After training, we can convert this SNN to Lava DL and check the accuracy:

net_ladl = net.to_lava().to(args.device)
net_ladl.eval()
test_loss = 0
test_acc = 0
test_samples = 0
with torch.no_grad():
    for img, label in test_data_loader:
        img = img.to(args.device)
        label = label.to(args.device)
        img = img.unsqueeze(0).repeat(args.T, 1, 1, 1, 1)
        img = encoder(img)
        img = lava_exchange.TNX_to_NXT(img)
        fr = net_ladl(img).mean(-1)
        loss = F.cross_entropy(fr, label)

        test_samples += label.numel()
        test_loss += loss.item() * label.numel()
        test_acc += (fr.argmax(1) == label).float().sum().item()

test_loss /= test_samples
test_acc /= test_samples

print('test acc[lava dl] =', test_acc)

Finally, we can export the SNN in Lava DL format to an hdf5 file, which can then be read by Lava. Lava can rebuild the SNN and run the SNN on Loihi, or the CPU-simulated Loihi.Refer to Network Exchange (NetX) Library for more details.

The export function is:

def export_hdf5(net, filename):
    # network export to hdf5 format
    h = h5py.File(filename, 'w')
    layer = h.create_group('layer')
    for i, b in enumerate(net):
        handle = layer.create_group(f'{i}')
        b.export_hdf5(handle)

export_hdf5(net_ladl, os.path.join(args.out_dir, 'net_la.net'))

The complete codes are stored in spikingjelly.activation_based.examples.lava_mnist. The arguments are defined as:

(lava-env) wfang@mlg-ThinkStation-P920:~/tempdir/w1$ python -m spikingjelly.activation_based.examples.lava_mnist -h
usage: lava_mnist.py [-h] [-T T] [-b B] [-device DEVICE] [-data-dir DATA_DIR]
                    [-channels CHANNELS] [-epochs EPOCHS] [-lr LR] [-out-dir OUT_DIR]

options:
-h, --help          show this help message and exit
-T T                simulating time-steps
-b B                batch size
-device DEVICE      device
-data-dir DATA_DIR  root dir of the MNIST dataset
-channels CHANNELS  channels of CSNN
-epochs EPOCHS      training epochs
-lr LR              learning rate
-out-dir OUT_DIR    path for saving weights

When we run this script, it will firstly train a SNN, then convert the SNN to Lava DL format and run an inference, and finally export the SNN to the hdf5 file:

(lava-env) wfang@mlg-ThinkStation-P920:~/tempdir/w1$ python -m spikingjelly.activation_based.examples.lava_mnist -T 32 -device cuda:0 -b 128 -epochs 16 -data-dir /datasets/MNIST/ -lr 0.1 -channels 16
Namespace(T=32, b=128, device='cuda:0', data_dir='/datasets/MNIST/', channels=16, epochs=16, lr=0.1, out_dir='./')
Namespace(T=32, b=128, device='cuda:0', data_dir='/datasets/MNIST/', channels=16, epochs=16, lr=0.1, out_dir='./')
epoch = 0, train_loss = 1.7607, train_acc = 0.7245, test_loss = 1.5243, test_acc = 0.9443, max_test_acc = 0.9443

# ...

Namespace(T=32, b=128, device='cuda:0', data_dir='/datasets/MNIST/', channels=16, epochs=16, lr=0.1, out_dir='./')
epoch = 15, train_loss = 1.4743, train_acc = 0.9881, test_loss = 1.4760, test_acc = 0.9855, max_test_acc = 0.9860
finish training
test acc[sj] = 0.9855
test acc[lava dl] = 0.9863
save net.state_dict() to ./net.pt
save net_ladl.state_dict() to ./net_ladl.pt
export net_ladl to ./net_la.net