Converting CSP Networks from Keras to Nengo

Hello everyone!
I am currently trying to convert a CSPNet style network from Keras to native Nengo. This means at some point I need to split my computation graph and slice the tensors like in this Keras code:

x1 = Conv2D(...)(inp)
# Slice tensor in half
x2 = x1[...,:x1.shape[-1]//2]
x2 = Conv2D(...)(x2)
# Concat the two tensors and fuse the graph
x1 = Concatenate(axis=-1)([x1, x2])
x1 = Conv2D(...)(x1)
model = Model(inp, x1)

The problem with this is that Keras internally uses the SlicingOpLambda to wrap the slice operation, which is (obviously) not implemented in the file. So naturally I looked into implementing my own custom LayerConverter class for converting this operation to Nengo.

This should be pretty simple since the only thing to do is to create two nodes with the second one being half the size of the preceding one and then connecting half of the preceding neurons. The code for this looks as follows:

class ConvertSlicingOpLambda(nengo_dl.converter.LayerConverter):
    def convert(self, node_id):
        output = self.add_nengo_obj(node_id)
        self.add_connection(node_id=node_id, obj=output, pre_slice=slice(output.size_out))
        return output

Note the pre_slice argument for the add_connection method. It is not there in the original method, with the reason being that the normal add_connection method doesn’t support connecting only part of the preceding node to the object given with obj. That is why I modified the original method like below:

def add_connection(self, node_id, obj, pre_slice=None, input_idx=0, trainable=False, **kwargs):
   if pre_slice is not None:
        # Only connect the part given by the pre_slice argument
        conn = nengo.Connection(
        # Just like before

With this “hack” my model does not perform the same in native Nengo as in Keras though, meaning something did not quite go right.

Is there any other method of connecting only part of a previous node with the one created in the convert method of a LayerConverter?

Thanks in advance, any help is much appreciated!

Just in case it is helpful, here is a full minimum example for showing the issue. Use the use_slice variable to switch between using or not using the custom operation.

import tensorflow as tf
import nengo_dl
import numpy as np
from tensorflow.python.keras.layers.core import SlicingOpLambda

physical_devices = tf.config.experimental.list_physical_devices("GPU")
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
    print("Called tf.config.experimental.set_memory_growth(GPU0, True)")

def evaluate_accuracy(predictions, ground_truth):
    prediction_values = np.argmax(tf.nn.softmax(predictions), axis=-1)
    ground_truth_values = np.argmax(ground_truth[:len(predictions)], -1)
    return (prediction_values == ground_truth_values).mean()

class ConvertSlicingOpLambda(nengo_dl.converter.LayerConverter):
    def convert(self, node_id):
        output = self.add_nengo_obj(node_id)
        self.add_connection(node_id=node_id, obj=output, pre_slice=slice(output.size_out))
        return output

# Parameter for comparing
use_slice = True

# Load MNIST dataset
dataset = tf.keras.datasets.mnist.load_data(path="mnist.npz")

inp = tf.keras.Input((28, 28, 1))
x1 = tf.keras.layers.Conv2D(16, 3)(inp)
if use_slice:
    x2 = x1[...,:x1.shape[-1]//2]
    x2 = x1
x2 = tf.keras.layers.Conv2D(16, 3, padding="same")(x2)
x1 = tf.keras.layers.Concatenate(axis=-1)([x1, x2])
x1 = tf.keras.layers.Conv2D(16, 1)(x1)
x1 = tf.keras.layers.Flatten()(x1)
x1 = tf.keras.layers.Dense(10)(x1)
model = tf.keras.Model(inp, x1)

# Train with keras
    optimizer='adam', metrics=['accuracy'])[0][0], tf.keras.utils.to_categorical(
    dataset[0][1], num_classes=10), batch_size=32, epochs=1)
print("Accuracy in second simulator: {}".format(
        tf.keras.utils.to_categorical(dataset[1][1][:100], num_classes=10)

# Reshape test data
test_data = dataset[1][0].reshape((dataset[1][0].shape[0], 1, -1))
test_target = np.expand_dims(
    tf.keras.utils.to_categorical(dataset[1][1], num_classes=10), 1)
# Create converter
converter = nengo_dl.Converter(model, allow_fallback=False)
nengo_inp = converter.inputs[inp]
nengo_outp = converter.outputs[x1]
# Simulator for testing with Nengo
with nengo_dl.Simulator(, minibatch_size=1) as sim:
    result = sim.predict({nengo_inp: test_data[:100]})[nengo_outp]
    print("Accuracy in second simulator: {}".format(
        evaluate_accuracy(result, test_target[:100])

Hi @MojoForPresident, and welcome to the Nengo forums. :smiley:

Your question is a tricky one for sure! I’m not a NengoDL power user by any means, and our lead NengoDL dev is on leave, so I couldn’t find a way to implement the slicing operator within the NengoDL converter. @drasmuss (the lead NengoDL dev) can weigh in on the converter approach when he returns.

However, playing around with your minimum working example, it looks like an alternative (working) approach is to re-implement the Keras model directly in NengoDL using the nengo_dl.Layer objects (like in this example)

There are some things I learned while doing this though:

  • When using the Conv2D layers, you need to specify the shape_in parameter. The NengoDL builder doesn’t seem to be able to pick it up automatically.
  • The nengo_dl.Layer class doesn’t yet handle multi-inputs and output Keras layers yet. Thus, for the tf.keras.layers.Concatenate layer, you need to wrap that within a nengo_dl.TensorNode (see this NengoDL Github issue)
  • The output of a nengo_dl.Layer is flat, so some extra math needs to be done to compute the appropriate indices you need to use for the slice. My approach was to create an appropriately shaped np.arange that I sliced, then flattened (It’s probably possible to make this cleaner, but I’ll leave that up to you. :smiley:):
x1 = nengo_dl.Layer(tf.keras.layers.Conv2D(16, 3))(inp, shape_in=(28, 28, 1))

ind = np.arange(x1.size_out).reshape(x1.shape_out)
ind = ind[..., : x1.shape_out[-1] // 2]
ind = ind.flatten()

x2 = x1[ind]

I took your minimum working example and implemented the nengo_dl.Layer version of it, while trying to mirror the Keras model (and model compilation, fitting and evaluation steps) as much as possible.
Here’s the code: (5.1 KB)

Assuming I have translated the Keras model to the nengo_dl.Layer format correctly, the model seems to train correctly now. Here’s an example run:

> python .\
Called tf.config.experimental.set_memory_growth(GPU0, True)
1200/1200 [==============================] - 6s 3ms/step - loss: 6.5578 - accuracy: 0.8105
4/4 [==============================] - 0s 9ms/step - loss: 0.4783 - accuracy: 0.9000
Accuracy in Keras simulator: [0.47832900285720825, 0.8999999761581421]

Build finished in 0:00:00
Optimization finished in 0:00:00
Construction finished in 0:00:00
1200/1200 [==============================] - 9s 6ms/step - loss: 5.2230 - probe_loss: 5.2230 - probe_accuracy: 0.8183
2/2 [==============================] - 0s 4ms/step - loss: 0.3524 - probe_loss: 0.3524 - probe_accuracy: 0.9100
Accuracy in NengoDL simulator: {'loss': 0.3524181842803955, 'probe_loss': 0.3524181842803955, 'probe_accuracy': 0.9100000262260437}
1 Like

Hey @xchoo, thank you so much for your more than detailed answer!
Your index “magic” with rearranging the np.arange index array actually solved my problem. The issue was that I naively connected the first half of the previous node to the following one:

self.add_connection(node_id=node_id, obj=output, pre_slice=slice(output.size_out))

But as you described:

As soon as I replaced my index logic with yours, my conversion worked!

class ConvertSlicingOpLambda(nengo_dl.converter.LayerConverter):
    def convert(self, node_id):
        # New index logic
        input_shape = self.input_shape(node_id)
        input_size =
        idx = np.arange(input_size).reshape(input_shape)
        idx = idx[..., :input_shape[-1]//2]
        idx = idx.flatten()

        output = self.add_nengo_obj(node_id)
        self.add_connection(node_id=node_id, obj=output, pre_slice=idx)
        return output

Thanks a lot again for your help! This forum is awesome :smile:

Huh. I didn’t realize that myself. haha. That’s great news! :smile: