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

xchoo · April 6, 2021, 6:42pm

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)

zerone · April 8, 2021, 6:25pm

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.

xchoo · April 14, 2021, 3:16am

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.

zerone · April 14, 2021, 4:31pm

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:

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)?
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!

xchoo · April 22, 2021, 1:34am

Hi @zerone,

This post is a little on the long side, and I’ve only just got the time to go through the entire post. To address your questions:

No, the synapse parameter affects all connections in the network, be it to an Ensemble or Node. However, I did misread your code, and you are correct in the observation that you do not see any filtered spikes. This is because the filtering is done on the output of a connection, which is then accumulated on the input of the Ensemble / Node. Thus, for nengo.Ensemble objects, to get the filtered inputs, you’ll need to probe the input attribute of the receiving Ensemble. As an example:

nengo.Probe(ndl_model.layers[conv2], attr="input")

For nengo.Node objects, there isn’t an input attribute that you can probe, so if you want to get the filtered spikes, you’ll need to probe the output of the connection that connects to that specific node (you’ll need to go through the all_connections object and find the one where post_obj == node)

Yes. As I mentioned above, you’ll need to probe the connection that connects to that node. Probably something like this (caveat: I haven’t tested this code)

for conn in ndl_model.net._connections:
    if conn.post_obj == ndl_model.layers[max_pool]:
        break
nengo.Probe(conn)

Note that this only works since the max_pool layer has 1 connection to it. If you make multiple connections to this node, you’ll need to probe all of the connections, then sum up the probed results for each of the probe.

Hmm, yes, I admit this is unexpected. Looking at the code some more, I believe this is a bug, and I have made a GitHub issue (along with the suggested fix).

This is indeed the case in the code.

This is correct. See this thread (where you posed the question before. ). With max pooling, you’ll definitely want to use some kind of averaging or smoothing to the spiking output to better approximate a rate-based max pooling operation.

It would be the same for all layers, and by that I mean that the filtered spikes is used as an input to each layer (regardless of where it is in the network). For the first conv0 layer, the filter is None (i.e, no filter is applied), and the “spike train” is just a steady value, but the math still applies.

zerone · April 22, 2021, 2:56pm

Hello @xchoo, NP and thanks for getting back.

With respect to the following:

you previously in another post (I don’t remember which) noted that filtering is done while inputing to the Ensembles (and not at the output of the Ensemble), however, I guess, this statement is equivalent to

i.e. output of a connection is exactly what’s input to the Ensemble/Node (with no other in-between processing of values).

With respect to the following,

I haven’t tested it yet, but just confirming with you… it still won’t work for the TensorNode in question right? as you realize it as a bug in current Nengo-DL version… isn’t it?

However, with respect to the bug you raised, if one still wishes to work on the spike amplitude output (and not the filtered one), then I guess he/she has to explicitly set the synapse of Connection to the respective Node/Ensemble/TensorNode as None. Isn’t it?

With respect to the following,

it means that ndl_model.layers and ndl_model.net.ensembles are essentially representing the same thing… right?

Yeh , thanks for mentioning the following… I should try it as well.

With respect to the following,

just explicitly confirming, the filtered “spike train” is considered as x at all the layers and current J is calculated for each of them. Isn’t it?

zerone · April 22, 2021, 5:35pm

With respect to the bug you raised, I noticed this right now. For the following model:

# 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)

# Default pool_size = (2,2), padding = "valid", data_format = "channels_last".
max_pool0 = tf.keras.layers.MaxPool2D()(conv0) 

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

max_pool1 = tf.keras.layers.MaxPool2D()(conv1) 

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


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

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

following is the output w.r.t. connection synapse:

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

for conn in ndl_model.net.all_connections:
  print(conn, "| Synapse: ", conn.synapse)

Output:

<Connection from <Node "conv2d.0.bias"> to <Node "conv2d.0.bias_relay">> | Synapse:  None
<Connection from <Node "conv2d.0.bias_relay"> to <Neurons of <Ensemble "conv2d.0">>> | Synapse:  None
<Connection from <Node "input_1"> to <Neurons of <Ensemble "conv2d.0">>> | Synapse:  None
<Connection from <Neurons of <Ensemble "conv2d.0">> to <TensorNode "max_pooling2d">> | Synapse:  None
<Connection from <Node "conv2d_1.0.bias"> to <Node "conv2d_1.0.bias_relay">> | Synapse:  None
<Connection from <Node "conv2d_1.0.bias_relay"> to <Neurons of <Ensemble "conv2d_1.0">>> | Synapse:  None
<Connection from <TensorNode "max_pooling2d"> to <Neurons of <Ensemble "conv2d_1.0">>> | Synapse:  None
<Connection from <Neurons of <Ensemble "conv2d_1.0">> to <TensorNode "max_pooling2d_1">> | Synapse:  None
<Connection from <Node "conv2d_2.0.bias"> to <Node "conv2d_2.0.bias_relay">> | Synapse:  None
<Connection from <Node "conv2d_2.0.bias_relay"> to <Neurons of <Ensemble "conv2d_2.0">>> | Synapse:  None
<Connection from <TensorNode "max_pooling2d_1"> to <Neurons of <Ensemble "conv2d_2.0">>> | Synapse:  None
<Connection from <Node "dense.0.bias"> to <TensorNode "dense.0">> | Synapse:  None
<Connection from <Neurons of <Ensemble "conv2d_2.0">> to <TensorNode "dense.0">> | Synapse:  Lowpass(tau=0.005)

As you can see, only the last connection from 3rd Conv layer to Output dense layer has a connection of 0.005.

zerone · April 22, 2021, 11:07pm

Sorry for so many questions… I cannot promise if this is my last one… .

So I have a Nengo network of K neurons which requires N input values and outputs one value each timestep. The N inputs are supposed to be the outputs of N neurons. I want to put this network between two Ensembles, but it comes with a catch. Just like the MaxPool op, each of the N inputs to this network from the previous Ensemble should be grid shaped and the single outputs should also be arranged in a grid (such that when flattened, it can be inputted to the next Ensemble or can have direct connection to the next Ensemble neurons).

Therefore, taking the pictorial example above, the 4 x 4 square represents the Ensemble neurons arranged in grid, there will be 4 instances of my Nengo network (each instance having its own K neurons) and each of the N neurons from the coloured sub-grid will be connected to an instance, with the 4 outputs forming the grid which will be directly input to the next Ensemble of 4 neurons (one-to-one connection).

The desired architecture can be seen in the picture below.

All the blue coloured Nengo Nets are independently functional and receive input in grid format, and their output is directly connected to the next Ensemble. Of course each of the inputs to the blue Nengo Nets are neuron’s input and their (i.e. Nengo Nets) output is fed to the individual neurons in next 2x2 Ensemble.

Keeping in mind the number of filters in a Conv Ensemble (previous and next both), how do I code this architecture? The reason I want such an architecture is to allow the neurons in the Nengo Nets to maintain their (voltage) state in each timestep of the simulation.

xchoo · April 28, 2021, 3:06am

Yes. This is expected (with the bug). This is because the connections are made by the post layer. So, when the model is processed, it goes something like this:

Add input layer, no connections made
Add conv0 layer. Connection from input layer made, synapse=None because input layer are not neurons.
Add max_pool0 layer. Connection from conv0 layer made, synapse=None (bug!). If correctly coded, the synapse=0.005 instead of None.
Add conv1 layer. Connection from max_pool0 layer made, synapse=None because preceding layer does not contain neurons (it’s a TensorNode)
Repeat for remaining max_pool and conv layers.
Add dense layer. This is combined with the flatten layer because the flatten layer is just a linear transform. Connection is made to conv2, and synapse=0.005 because conv2 contains neurons.

I haven’t implemented this myself, but I think the easiest way to do this is to use the nengo.networks.EnsembleArray. The EnsembleArray is essentially a bunch of ensembles that are grouped into a single network. The convenience of the EnsembleArray comes with working with the input and outputs. With the EnsembleArray, the entire network is treated like one big ensemble, so all you have to do is make one connection to it, and internally it will handle the projections from the input dimensions to the individual ensembles in the EnsembleArray.

In summary, create an EnsembleArray of (in the case of the picture you posted) 4 sub-ensembles. Each sub-ensemble should be set to a dimensionality (using ens_dimensions) of 4 since they will each be representing 4 values. Finally, connection from the previous_ens to the EnsembleArray and define the connection weight matrix (i.e. the transform value in the nengo.Connection) in such a way as to implement the mapping you want. The transform value should be a programmatically shuffled identity matrix (I’ll leave it up to you to figure out how to do this).

zerone · April 28, 2021, 4:05am

Hello @xchoo, thanks for the suggestion. Actually, I had in mind that EnsembleArray might be useful in my case. However, I was stuck at two issues,

(1): how to input from the previous Ensemble, such that the neuron mapping is maintained? → I guess, this can be done by transform parameter to map the values (I am yet to realize it programmatically).

(2): My individual Nets which receive input from the 4 neurons are not just an ensemble of K neurons, but a specially designed Nengo network of passthrough Nodes and Ensembles. I also want the dynamics of the neurons in this network to evolve with that of the main converted network (which I guess, EnsembleArray takes care of). Therefore I am stuck at how to repeat this specially designed network over the grid of neurons.

In a nutshell, is there a way in EnsembleArray such that I can mention an instance of this Net as a parameter and it gets repeated? If not, then am I left with the only option of individually connecting the neurons of the previous Ensemble to my Net in a loop? Please let me know.

EDIT:
I was looking for more options w.r.t. above and came across Designing networks tutorial. So looks like, I can design my Net in a function and call the function again and again to create new instances of it. Next, all is left is to appropriately connect the different instances of my Net to the corresponding neurons in the grid layout of pre_ensemble and post_ensemble. Now, these Ensembles receive values and output values in flattened vectors shape. So I guess, I need to do some simple maths to map a tuple of (row_index, col_index, channel_index) of the implicit grid layout to the corresponding vector index in the flattened vector for both the pre_ensemble and post_ensemble while inputting to and outputting from my Net. Before I implement it, I wanted to know if there’s a better way than this? Also, if it sounds relevant to you, can you please elaborate on the following

How should the transform look like for an input of say… (26, 26, 32) shaped matrix? I did not get the idea of shuffled identity matrix. If you can explain it, I should be able to figure out how to create it. Thanks!

xchoo · April 29, 2021, 4:46am

I was thinking about it some more, and I think you can also do the connection using the numpy slicing operator.
As an example, consider your input is a 2D array like so:

 0  1  2  3 
 4  5  6  7
 8  9 10 11
12 13 14 15

If you want to do a max pooling operation with a 2x2 kernel, you’d want elements 0, 1, 4, and 5 to go to the first ensemble, then elements 2, 3, 6, and 7 to go to the second ensemble, and so forth. If the input to the ensemble array (or a custom network) is 16D, the numy slice would look like this:

0 1 4 5 2 3 6 7 8 9 12 13 10 11 14 15

You can use that slice when you make the nengo.Connection:

slice = [0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15]
nengo.Connection(ens.neurons[slice], max_pool.input)

Here’s some code that will generate the slice for you (this is hardcoded for a 4x4 input with a 2x2 kernel):

    trfm = np.zeros(16, dtype=int)
    mat = np.arange(16).reshape((4, 4))
    n_kernels = int(16 / 4)
    for n in range(n_kernels):
        trfm[n * 4 : n * 4 + 4] = mat[
            n // 2 * 2 : n // 2 * 2 + 2, n % 2 * 2 : n % 2 * 2 + 2
        ].flatten()

It’s not the prettiest, and I think there are better ways of doing it, but I have tested it to work.

You can do this, yes. You can even make a network of networks. I.e., make a “main” network that serves as the interface to the rest of your model, and within that network, call the network creation function for the subnetworks.

zerone · April 29, 2021, 4:00pm

Thank you @xchoo for the slicing suggestion and directly connecting the sliced neurons to the max_pool.input object. This very much cleans the code.

However, one needs to have a max_pool Nengo object (which I believe should be a Network here) which I have little idea to create. I have seen examples of it, like in spa networks where you connect an Ensemble to spa.<network>.input, but how do I structure such a “main” network, such that the sliced inputs get distributed in groups of 4s to each sub network? An example to create such will help. Or w.r.t. the following:

is it straight ahead to simply call the subnetwork creation function multiple times in a for loop (within a “main” network) and connect it to the “main” network’s input - but how to connect each subnetwork such that each subnetwork accepts sliced inputs in groups of 4s?

xchoo · April 30, 2021, 3:05am

You can use the same slicing trick to distribute an N-dimensional input into groups of 4. As an example:

with nengo.Network() as max_pool_main:
    # Note, you can use a function to define your function here as well, 
    # this is just an example
    max_pool_main.input = nengo.Node(size_in=16)
    max_pool_main.output = nengo.Node(size_in=4)
    
    for n in range(16 // 4):
        # Make individual subnets and connect to input
        subnet = create_subnet()
        # Connect sliced input to input of subnet
        nengo.Connection(max_pool_main.input[n * 4: n * 4 + 4], subnet)
        # Connect output of subnet to sliced output
        nengo.Connection(subnet, max_pool_main.output[n])

Note that in the example code above, it makes no assumptions about the order of each element of the input (and it doesn’t need to). All it is doing is “mapping” the groups of 4 dimensions of the input to the various copies of subnet. You can combine it with the shuffled slicing I suggested in my previous post to get the correct mapping for the max pooling operation, like so:

slice = [0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15]
nengo.Connection(ens.neurons[slice], max_pool_main.input)

If you are looking for an example network to start, I’d suggest the nengo.networks.EnsembleArray network code. It basically have everything you want (just use the __init__ function, and you can ignore the rest really), all you have to do to customize it to your needs is to replace this line with code to create your subnet network.

zerone · April 30, 2021, 5:08pm

Thank you very much @xchoo for giving an example of it as well as the reference to EnsembleArray code. I will first start with the basics to get it correct, and then look for automatic replacement of the layers by registering it to the Converter. I guess, EnsembleArray code would be useful then. Until then, can you please let me know your thoughts over these questions in post #7. I guess it was missed in the plethora of my questions.

xchoo · May 16, 2021, 1:37am

Hi @Zerone
I didn’t realize that I had missed some of your questions:

Probing the output of the connection that connects to a TensorNode is the only way to get the input signal to the TensorNode. It will still work with the bug in the current version of NengoDL. The bug only affects the synaptic filtering on that connection, so you will get a spiky output regardless of what synapse is set in the converter.

For this bug, the synapse for connections to TensorNodes is always None. So, if you wanted to work with the spike output, you wouldn’t need to do anything.

Not quite. ndl_model.layers is a dictionary (i.e., a map) between the TensorFlow layer object and the corresponding Nengo ensemble.

The filtered spike train weighted by the input connection weight matrix yes.

zerone · May 16, 2021, 4:49pm

Hello @xchoo, no problem and thanks for getting back!

Exactly my thoughts. This is what I wanted to confirm, I think I should have rephrased my question well.

With respect to the following,

I meant to ask… after the bug is fixed and I want to work with the spike output, then I would have to set the synapse=None explicitly in the Connection to TensorNode before simulating the network… right? And thanks for confirming the following!

BTW, when I implemented your suggestion with slight modifications to suit my needs, and simulated the network; it takes nearly:

|    #            Optimizing graph: creating signals                  | 0:24:37
|                          Optimizing graph                           | 0:39:17
Optimizing graph: creating signals finished in 0:24:37
|                          Optimizing graph                           | 0:39:17
Optimization finished in 0:39:17
|#                        Constructing graph                          | 0:00:00

Once again for reference, my create_subnet() returns a subnet composed of PassThrough nodes and Ensembles, and the neurons in the Ensemble are directly connected to the nodes (i.e. with nengo.Connection(ens.neurons[i], node, ...)), so I guess, there’s no computation of decoders happening, yet it takes significant amount of time to Optimize the network (even more than an hour in case of few wider nets). After optimization, the network functionally works as expected, but was wondering if you would have any suggestions to bypass/improve the optimization time? Please let me know!

xchoo · May 26, 2021, 3:08am

You may want to try disabling the operator merging. This might improve the optimization time.

zerone · May 26, 2021, 9:56pm

Thank you @xchoo for the suggestion. I am bit caught up with another work, but will try it soon and let you know.

zerone · June 16, 2021, 3:58am

Hello @xchoo, I have few questions below:

1> In a nengo.Connection(), what is the relative ordering of synapse, function, and transform? I know that function is applied before the transform, so is it synapse → function → transform, and in absence of function, is it synapse → transform? i.e. incoming spikes (if the pre_obj is an Ensemble) are synapsed/filtered first and then a Linear Transform is applied or is it just the opposite?

2> The default spike amplitude is 1/dt (due to the amplitude default set to 1). When I probe the Ensemble neurons of type nengo.LIF() or nengo.SpikingRectifiedLinear() or nengo_loihi.neurons.LoihiSpikingRectifiedLinear() (with default values of amplitude and dt), I see the spikes of amplitude 1000.0 (2000.0 when amplitude is set to 2). In NengoDL, the spikes are calculated as sim.data[neurons_probe] x scale_firing_rates x dt, which based on above (assuming amplitude=1, scale_firing_rate=100, dt=0.001) would evaluate to 1000.0 x 100 x 0.001 = 100, but based on my experience with NengoDL, for a scale_firing_rate of 100 (or 250), I see the spikes (calculated by the same previous formula) as 10 (or 4). Why is it not equal to 100 (or 250) (as should be obtained by the formula)? It seems like we should rather divide the amplitude with scale_firing_rates (as mentioned in your excellent description here). For some reason, I am messing things up. Please clarify? What computations do the spikes of amplitude 1/dt go through after being generated from the previous ensemble neurons in NengoDL?

3> This might sound trivial, but just confirming: SpikingRectifiedLinear() and LoihiSpikingRectifiedLinear() neurons are examples of “Integrate and Fire” neurons… right?

xchoo · June 19, 2021, 3:07pm

In the build model (i.e., the model that the simulator uses), both the function and transform are combined into the connection’s weight matrix. Recall that the function is not a physical entity. It is just used by the connection’s weight solver to solve for the appropriate connection weights that will approximate the function.

As for the synapse, it is applied after the connection weight has been applied. However, since both the connection weight and synapse application are linear transforms, mathematically, it doesn’t matter which order they are done.

Can you provide a minimal network that demonstrates this? Just so that I have context with what network you are building, the parameters you have used, and how you are testing it.

This is correct.