[Nengo-DL]: Adding a dynamic custom layer between two Ensembles

Hello there,

So I want to add a dynamic layer (actually a nengo.Node containing my custom code) between two Ensembles after the network has been trained in TF mode and converted to spiking one. Precisely speaking, I intend to replace a particular type of TF layer (which will be a TensorNode after conversion) with my custom nengo.Node layer. Of course, the nengo.Node will work on the inputs from the previous nengo.Ensemble and output to the next nengo.Ensemble. Is it possible and if yes… how? I couldn’t find any related example. I understand that this question might appear vague, but please let me know if you need any more information.

Thanks!

Edit: I came across this article on extending the nengo_dl.Converter() but it seems that it applies to the cases when we intend to create a new custom layer before training process.

Since you are attempting to make modifications to the network after the network has been converted, what you’ll want to do is to modify the converter.net Nengo network (after you load any weights into it).

Modifying an existing Nengo network is something that is not streamlined in Nengo, but because Nengo is written in Python, it is not too difficult to implement. Each nengo.Network has a list for each Nengo object contained within it, i.e., it has lists _connections, _nodes, and _probes (and _ensembles if you use it). With this, and a bit of Python logic, we can figure out the things connected to the node we want to replace, and re-create that connection to the replacement node. We can also use these lists to remove references to the node we want to replace, so they aren’t actually run when the network runs.

Here’s how you would iterate through all of the connections in a network to find out the one that’s connected to a specific node:

with model:
    for conn in model.all_connections:
        if conn.post_obj is my_node:
            break

Once you find the connection in question, you can then create new connection with the same properties as the old connection, but connected to the replacement node:

    nengo.Connection(
        conn.pre_obj,
        replacement_node,
        transform=conn.transform,
        synapse=conn.synapse,
        function=conn.function,
    )

Finally, you can remove the old connection, the old node, and any probes attached to that node:

    model._connections.remove(conn)
    model._nodes.remove(conn.post_obj)
    for probe in model.all_probes:
        if probe.target is conn.post_obj:
            model._probes.remove(probe)

Here’s a quick script (not in NengoDL, just regular Nengo) demonstrating the whole process in one file: test_nengo_replace.py (1.8 KB)

Thank you @xchoo for looking into it. I have experimented with your suggested method of replacing a TensorNode and was seemingly successful. However, I have further questions and for the same, I am taking an example of simple 2D CNN network (taken from Nengo-DL examples) and added a MaxPooling op to result in creation of a TensorNode after conversion. Below is the TF network.

# input
inp = tf.keras.Input(shape=(28, 28, 1))

# convolutional layers
conv0 = tf.keras.layers.Conv2D(
    filters=32,
    kernel_size=3,
    activation=tf.nn.relu,
)(inp)

max_pool = tf.keras.layers.MaxPool2D()(conv0)

conv1 = tf.keras.layers.Conv2D(
    filters=64,
    kernel_size=3,
    strides=2,
    activation=tf.nn.relu,
)(max_pool)

# fully connected layer
flatten = tf.keras.layers.Flatten()(conv1)
dense = tf.keras.layers.Dense(units=10, activation="softmax")(flatten)

model = tf.keras.Model(inputs=inp, outputs=dense)

Its model.summary() output is below.

Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 28, 28, 1)]       0         
_________________________________________________________________
conv2d (Conv2D)              (None, 26, 26, 32)        320       
_________________________________________________________________
max_pooling2d (MaxPooling2D) (None, 13, 13, 32)        0         
_________________________________________________________________
conv2d_1 (Conv2D)            (None, 6, 6, 64)          18496     
_________________________________________________________________
flatten (Flatten)            (None, 2304)              0         
_________________________________________________________________
dense (Dense)                (None, 10)                23050     
=================================================================
Total params: 41,866
Trainable params: 41,866
Non-trainable params: 0
_________________________________________________________________

As can be seen above, output from the first conv2d layer is of shape (26, 26, 32) which is input to the max_pooling2d layer. The max-pooling op has a kernel of shape (2, 2) and strides by default is the same, thus the output from the max_pooling2d layer is of shape (13, 13, 32) i.e. image dimensions have been halved and number of filters/channels remain the same. The input to the next conv2d_1 layer (it has a stride of 2) is of the same shape i.e. (13, 13, 32) and thus the output is of shape (6, 6, 64).

Now, let’s try replacing this max_pooling2d TensorNode (after training and conversion). The expectation from the new custom Node would be to accept an input of shape (26, 26, 32) and output data of shape (13, 13, 32) which will be then an input to the next conv2d_1 layer. The operation in the custom Node should respect the topography of the input to it. But how to ensure that Node receives and preserves the correct input topography? Or in other words, how to do the operations in the custom Node which preserves the input/output topography? For the MaxPooling op shown below (source: wikipedia)


I not only have to preserve the above spatial input, but also the filter order (shown below)
MaxpoolSample

How do I preserve the correct input and output topology to/from the custom Node?

For now, I am replacing the max_pooling2d TensorNode with the following code.

def func(t, x):
  print(x.shape)
  return x[:5408]

with ndl_model.net:
  # Create Custom Node.
  new_node = nengo.Node(output=func, size_in=21632, label="Custome Node")

  conn_from_conv0_to_max_node = ndl_model.net.all_connections[3]

  # COnnection from Conv0 to MaxPool node.
  nengo.Connection(
    conn_from_conv0_to_max_node.pre_obj,
    new_node,
    transform=conn_from_conv0_to_max_node.transform,
    synapse=conn_from_conv0_to_max_node.synapse,
    function=conn_from_conv0_to_max_node.function)

  # Connection from MaxPool node to Conv1.
  conn_from_max_node_to_conv1 = ndl_model.net.all_connections[6]
  nengo.Connection(
    new_node,
    conn_from_max_node_to_conv1.post_obj,
    transform=conn_from_max_node_to_conv1.transform,
    synapse=conn_from_max_node_to_conv1.synapse,
    function=conn_from_max_node_to_conv1.function)

  # Remove the old connection to MaxPool node and from MaxPool node, MaxPool node.
  ndl_model.net._connections.remove(conn_from_conv0_to_max_node)
  ndl_model.net._connections.remove(conn_from_max_node_to_conv1)
  ndl_model.net._nodes.remove(conn_from_conv0_to_max_node.post_obj)

I also checked if making the new connection is successful.

ndl_model.net.all_connections
[<Connection at 0x2b97fffc01d0 from <Node "conv2d.0.bias"> to <Node "conv2d.0.bias_relay">>,
 <Connection at 0x2b97efeee190 from <Node "conv2d.0.bias_relay"> to <Neurons of <Ensemble "conv2d.0">>>,
 <Connection at 0x2b97fffbd990 from <Node "input_1"> to <Neurons of <Ensemble "conv2d.0">>>,
 <Connection at 0x2b97fff8b810 from <Node "conv2d_1.0.bias"> to <Node "conv2d_1.0.bias_relay">>,
 <Connection at 0x2b97efe79590 from <Node "conv2d_1.0.bias_relay"> to <Neurons of <Ensemble "conv2d_1.0">>>,
 <Connection at 0x2b97effc8690 from <Node "dense.0.bias"> to <TensorNode "dense.0">>,
 <Connection at 0x2b984b73fad0 from <Neurons of <Ensemble "conv2d_1.0">> to <TensorNode "dense.0">>,
 <Connection at 0x2b9b8a11b490 from <Neurons of <Ensemble "conv2d.0">> to <Node "Custome Node">>,
 <Connection at 0x2b9b8a12ab90 from <Node "Custome Node"> to <Neurons of <Ensemble "conv2d_1.0">>>]

and from the above output, it seems it is. Also, the output of print(x.shape) in the func() passed in the custom Node is (21632,) = 26 x 26 x 32 is coalesced in one vector. Interestingly, why is the order of all the connections not sequential?

I also probed the ndl_model.net first Conv layer with following code:

with ndl_model.net:
  # Output from the first Conv layer.
  # ndl_model.layers[conv0].probeable => ('output', 'input', 'output', 'voltage')
  conv0_lyr_otpt = nengo.Probe(ndl_model.layers[conv0], attr="output")
  # ndl_model.net.ensembles[0].probeable => ('decoded_output', 'input', 'scaled_encoders')
  conv0_ens_otpt = nengo.Probe(ndl_model.net.ensembles[0], attr="decoded_output")
  
  # ndl_model.layers[conv0].ensemble.neurons.probeable => ('output', 'input', 'output', 'voltage')
  conv0_lyr_nrns_otpt = nengo.Probe(ndl_model.layers[conv0].ensemble.neurons, attr="output")
  # ndl_model.net.ensembles[0].neurons.probeable => ('output', 'input', 'output', 'voltage')
  conv0_ens_nrns_otpt = nengo.Probe(ndl_model.net.ensembles[0].neurons, attr="output")

I was just curious about the output from ndl_model.layers and ndl_model.net.ensembles from the first conv2d layer and it seems that the outputs from both sources match, as shown below… so they are essentially representing the same thing… right?

neuron_index = 16345
print(data1[conv0_lyr_otpt][0, :, neuron_index])
print(data1[conv0_lyr_nrns_otpt][0, :, neuron_index])
print(data1[conv0_ens_nrns_otpt][0, :, neuron_index])
[0.       0.       0.       9.999999 0.       0.       0.       0.
 9.999999 0.       0.       0.       9.999999 0.       0.       0.
 0.       9.999999 0.       0.       0.       9.999999 0.       0.
 0.       0.       9.999999 0.       0.       0.       0.       9.999999
 0.       0.       0.       9.999999 0.       0.       0.       0.      ]
[0.       0.       0.       9.999999 0.       0.       0.       0.
 9.999999 0.       0.       0.       9.999999 0.       0.       0.
 0.       9.999999 0.       0.       0.       9.999999 0.       0.
 0.       0.       9.999999 0.       0.       0.       0.       9.999999
 0.       0.       0.       9.999999 0.       0.       0.       0.      ]
[0.       0.       0.       9.999999 0.       0.       0.       0.
 9.999999 0.       0.       0.       9.999999 0.       0.       0.
 0.       9.999999 0.       0.       0.       9.999999 0.       0.
 0.       0.       9.999999 0.       0.       0.       0.       9.999999
 0.       0.       0.       9.999999 0.       0.       0.       0.      ]

except that conv0_ens_otpt = nengo.Probe(ndl_model.net.ensembles[0], attr="decoded_output") produces a output of shape (n_test_images, n_steps, 1), i.e. perhaps the overall decoded output from the Ensemble of 21632 neurons. Right?

One interesting thing I noted from above output was that the inputs to the max_pooling2d layer are the scaled spikes (I thought the inputs should have been smoothed values), but I guess, it goes with what you mentioned earlier, the outputs from the Ensemble neurons are the spikes and the smoothing happens during the input to the next Ensemble, and since MaxPooling2D op is not an Ensemble it simply operates on scaled spikes every timestep e.g. max(9.9999, 0, 0, 0) = 9.9999 i.e. simply output the max amplitude spike as an input to the next Ensemble… right?

Sorry for such a long response, please let me know if anything is confusing here.

The input to a Nengo node is always a 1-dimensional (flat) array. So, whatever function you implement for that node will need to handle that flat array. As for input mapping, you can use np.flatten to flatten an example multi-dimensional array into a flat array to see how the mapping is achieved. Alternatively, you can j use the np.reshape function to reshape the flat array back into a multi-dimensional array, and then do the max-pooling operation on that. This stack overflow thread discusses some of the approaches to using numpy operations to do the max-pooling operation.

The connections in the connection list is ordered by order of creation. And from my quick glance at it, it seems to be respecting that.

I have to double check the code, but if I recall correctly, I believe this is the case. If your ndl_model is NengoDL converter object, it contains references to the Nengo equivalents to each Keras layers (i.e., using .layers), as well as direct references to the Nengo model objects (i.e., using .net.ensembles).

For NengoDL converted networks, I wouldn’t use the decoded_output of any ensembles. Remember that for NengoDL models, encoders and decoders don’t really make much sense because the NEF hasn’t been used to optimize the connection weights.

I think this depends on what you have set for the synapse parameter when you create the NengoDL converter. If you leave the synapse value at None, the NengoDL converter will make a non-filtered connection between the ensemble and your custom node. This would mean that your node would receive spiking input instead of filtered ones.

As a side note, I discuss some other options for implementing the MaxPool operator using a custom Layer builder class in this forum thread. Not sure if that would be helpful to you.

Hello @xchoo, thanks for the stack-overflow source, and for your suggestion to use np.reshape; I checked, it maintains the order of the matrix same as before, after flattening it.

With respect to the following,

yes, looks like it maintains the order of connection creation, I thought that it would maintain the actual hierarchical order of connections. And yes, I agree with the following, thanks for reminding!

With respect to the comment below:

sorry, but it doesn’t seem to be happening in my case. Here’s the Converter code I experimented with:

ndl_model = nengo_dl.Converter(model, 
                               swap_activations={tf.nn.relu: nengo.SpikingRectifiedLinear()},
                               scale_firing_rates=100,
                               synapse=0.005)

and as you can see, synapse is not None. I believe, this synapse parameter becomes relevant at the synaptic connection to the Ensembles, and not to the Nodes (as Nodes aren’t supposed to represent a group of neurons, thus no synaptic connection required… right?). Thus, changing the synapse value in Converter will have effect on the input to the Ensembles, i.e. filtering the spikes accordingly.

I ran my code right now, and it produces the same non-filtered spike amplitudes (as the output of the first conv0 layer which is the supposed input to the max_pool layer) with synapse = 0.005 in the Converter.

Unfortunately, ndl_model.layers[max_pool].probeable produces the output ('output',), thus I cannot Probe the “input” to the max_pool TensorNode, and have to rely on the output of the previous conv0 layer (probing it with a synapse None to check the actual output) as done by conv0_lyr_otpt = nengo.Probe(ndl_model.layers[conv0], attr="output") (default Probe synapse is None).

If there’s a way to probe the “input” to the max_pool TensorNode directly, please let me know. I can confirm that after modifying the connection though, I can see unfiltered spike amplitude input to the custom new_node as follows:

Code:

def func(t, x):
  print(x.shape)
  print(x[:200])
  return x[:5408]

Output at each time-step (of course the output below is not fixed, some time-steps it’s all zeros):

(21632,)
[0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       9.999999 0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.
 0.       0.       0.       0.       0.       0.       0.       0.      ]

Now, if we go by your proposition that mentioning a synapse value in the Converter will create a filtered input channel to the custom Node, and that it may get null when the connection is replaced (which is my proposition), then this too doesn’t seem to happening either; as can seen in the code below for creating a new connection from the conv0 layer to the custom new_node:

# COnnection from Conv0 to MaxPool node.
  nengo.Connection(
    conn_from_conv0_to_max_node.pre_obj,
    new_node,
    transform=conn_from_conv0_to_max_node.transform,
    synapse=conn_from_conv0_to_max_node.synapse,
    function=conn_from_conv0_to_max_node.function)

It is expected that transform, synapse, and function of the original connection be preserved in the new connection. Moreover, following is the output before replacing the original connection.

conn_from_conv0_to_max_node = ndl_model.net.all_connections[3]
print(conn_from_conv0_to_max_node.pre_obj)
print(conn_from_conv0_to_max_node.transform)
print(conn_from_conv0_to_max_node.synapse)
print(conn_from_conv0_to_max_node.function)

Output:

<Neurons of <Ensemble "conv2d.0">>
NoTransform(size_in=21632)
None
None

Hence, please check this too, along with the following

I have two extra small questions:

  1. If at all MaxPooling in Nengo-DL is simply done by taking a max of incoming spike amplitudes at each timestep, then isn’t this somewhat suboptimal? As, the TF MaxPooling op is supposed to choose the maximum activation among a group of neurons, but Nengo-DL MaxPooling (with TensorNode) chooses the maximum spike amplitude at each timestep and it is possible that at a certain timestep, the chosen spike amplitude might not correspond to the maximally spiking neuron (which represents the maximum activation at the end of a sufficiently long simulation time period)?

  2. We have discussed the following for a number of times now, that the current J = alpha x <e, x> + bias (where x is the repeated scalar input) is the input to the Ensembles, and I definitely see it being true for the first conv0 layer, but what about the next Ensembles after the first? As mentioned above in this reply that the input to the Ensembles is the filtered spike output from the previous layer, is the filtered spike output too converted to current J using the same above formula (i.e. the filtered spike output is considered x) , or is it just straight forward inputted to the next conv1 (and so on) layers?

Please let me know!