Adding a custom Memristor object to Nengo

Backends
1

I am working on incorporating some custom code for learning using memristors into the Nengo backend, as groundwork for implementation into NengoDL.

My focus is on efficiency, as my current code running in a Node really does not scale well with the pre and post number of neurons.
The reason for this is that I’m simulating a full matrix of memristors, one for each (pre-post) neuron pair, and updating this at each timestep, as can be seen in the following snippet:

    # create memristor array that implement the weights
    self.memristors = np.empty( (self.output_size, self.input_size), dtype=Memristor )
    for i in range( self.output_size ):
        for j in range( self.input_size ):
            self.memristors[ i, j ] = self.memristor_model( seed=seed, voltage_converter=self.voltage_converter,
                                                            base_voltage=self.base_voltage, gain=self.gain )
            self.weights[ i, j ] = self.memristors[ i, j ].get_state()

My custom learning rules check for spikes and updates the corresponding memristors:

    # we only need to update the weights for the neurons that spiked so we filter for their columns
    if spiked_map.any():
        for j, i in np.transpose( np.where( spiked_map ) ):
            update = self.memristors[ j, i ].pulse( signal[ j, i ] )
            # update = update if update >=
            self.weights[ j, i ] = update

I think I’m quite clear on the build process but I’m not sure of how much simply moving my code to a custom LearningRule could improve the running time, as using a Node and simulating 30 seconds with 100 pre/post neurons takes around one hour.

I was thus wondering if adding a new Memristor object type to Nengo could help. I would imitate the Neurons and Ensemble implementations to have an object representing a vector of memristors whose current resistances are kept track in a Signal, exactly how it is done for the voltage and refractory time of LIF neurons.
Would implementing this new Memristor object (that inherits from FrozenObject) in such a way help the efficiency? I was thinking that this way I could take advantage of the scheduler and optimiser in the Simulator.

Unfortunately, I am not sure of how to actually build such a custom object into the Simulator as it isn’t a Neuron, Ensemble, Connection, Synapse or Probe. I would have to somehow call model.build() outside of the Simulator but I can’t see a way to do that.
Is it only feasible to add custom implementations of Neuron, Ensemble, Connection, Synapse or Probe types?

2

Hi Tioz!

Adding a new Memristor object to Nengo may address the efficiency issues you are encountering with your code, however, adding a proper Memristor object to Nengo will involve delving deeply into the Nengo signals and objects builder logic, and that can be very time consuming to understand and implement properly.

Instead, (as a first pass, and quicker approach) you can try to optimize your code by utilizing the broadcasting features of numpy. Numpy has several built in functions that are optimized to perform operations on vectors and matrices much more efficiently that using a for loop, and is used extensively in Nengo (e.g., in the LIF neuron step_math method: https://github.com/nengo/nengo/blob/master/nengo/neurons.py#L493).

A quick example of this broadcasting feature:
Instead of doing this:

for i in range(10000):
    output[i] = math.sin(2 * i)

with Numpy, this can be done:

i = np.arange(10000)
output = np.sin(2 * i)
3

Hi @xchoo, thanks for your answer!

I had started implementing a Memristor object ( frontend and backend ) but got stuck on how to actually insert it into the build process… maybe it’s not a priority for me now, if I can increment the efficiency of the learning rule, but I would still be curious to know how I would go about building such a custom object into nengo.

Regarding using NumPy, I am well aware of the efficiency gains of using pure vectorialised code. Unfortunately the new resistance is a function of the old one, so I think I need to access each element separately. I could certainly parallelise the code so that each weight is updated separately, but I don’t see a way to vectorize it, but please correct me if I’m wrong!
I think that using the connection weight Signal to store the resistance and reading from it in my learning rule, could help somewhat. I would absolutely be happy for any suggestion on how to efficiently implement an array of memristors! I hate those for loops too :slight_smile: .

4

I took a look at your code, and you seem to have a good start to integrating the Memristor object into the Nengo build process. I haven’t actually tried to run your code (I’m not familiar with how Memristors work, or how to use them), so if you have a specific question or error about the integration process, please post them here.

As for increasing the efficiency of your code, one major difference I noticed between your code and Nengo’s is that pretty much all of Nengo performs operations on arrays of things (neurons, synapses, etc.) However, in your code, you are operating on individual Memristors (e.g., to do the pulse function).

The first step to vectorizing your code would be to modify your MemristorType class to store the parameters for a group of memristors, rather than just a single one (see the Nengo neurons.py file for an example of this: https://github.com/nengo/nengo/blob/414915af73c0284a1966967aecd98538c7fa33eb/nengo/neurons.py#L74). Once that has been done, the compute_pulse_number and compute_resistance functions can be vectorized using the appropriate numpy functions.

5

I didn’t not get a particular error, but was unclear on how to integrate a completely new NengoObject into the standard build process.
It is easy to add new Objects that extend pre-existing types by extending their xxType class, but the builder does not obviously integrate the calls to the build functions of new ObjectTypes.
So I can instantiate a BidirectionalPowerlaw() class but the Nengo builder will not call the build_bidirectionalpowerlaw() function and I don’t see an obvious way to call this function myself, because the model parameter is local to the Simulator, which in turn calls the nengo builder.

Regarding the efficiency of my code, I think something like this is much better:

        # analytical derivative of pulse number
        def deriv( n, a ):
            return a * n**(a - 1)
        
        # set update direction and magnitude (unused with powerlaw memristor equations)
        V = np.sign( pes_delta ) * 1e-1
        # use the correct equation for the bidirectional powerlaw memristor update
        # I am using forward difference 1st order approximation to calculate the update delta
        try:
            delta[ V > 0 ] = r_max * deriv( (((weights[ V > 0 ] - r_min) / r_max)**(1 / a)), a )
            delta[ V < 0 ] = -r_3 * deriv( (((r_3 - weights[ V < 0 ]) / r_3)**(1 / c)), c )
        except:
            pass

One thing I’m not completely clear about is the behaviour of the delta Signal. I would have expected the Signal to have a Reset() Operator applied to it at every timestep, but it instead seems to be incrementally updated.
Is that correct? Should we not be starting with a “clean” delta at each learning rule step?

6

*Marking this as the solution, but note that this is a multi-part question so there will be more answers in the discussion between Tioz and xchoo

As long as you register the Object with the Nengo builder, it should build it using the registered build function if the builder encounters it in your model. For your code in particular, I noticed that the build function for BidirectionalPowerlaw is in a different file than the class definition. I would recommend you put it in the same file so that importing the BidirectionalPowerlaw class should trigger the build function registration.

Yeah, that’s the general idea. But you’ll have to apply the same technique to the rest of your code as well. You’ll want to perform as many matrix/array type operations as possible to maximize the optimizations of numpy.

I’m not exactly sure which code you are referring to, but I presume it’s this (from https://github.com/nengo/nengo/blob/414915af73c0284a1966967aecd98538c7fa33eb/nengo/builder/learning_rules.py#L532)?

    delta = Signal(shape=target.shape, name="Delta")

    model.add_op(Copy(delta, target, inc=True, tag=tag))
    model.sig[rule]["delta"] = delta

In which case, this is correct. In Nengo, the Signal class encapsulates any values being transferred around to the various operators. In this case,

delta = Signal(shape=target.shape, name="Delta")

creates the Signal container object, then

model.add_op(Copy(delta, target, inc=True, tag=tag))

adds an the Copy operator to the model’s operator graph. With the inc=True flag set, what this operator is doing is taking the value in delta, adding it to target, and then copying that result back into target. Essentially: target = target + delta. No reset operator is needed here because the delta signal gets its value from an external source (external to the learning rule, at least). This line:

model.sig[rule]["delta"] = delta

then stores the delta Signal object for other learning rule subclasses to use in their respective operators.

It should be noted that if your learning rule is not of the type where the weights are modified with
weights = weights + delta, then you will want to create a new learning rule object specifically for your use case.

One important thing to understand is that subclassing a Nengo object (be it a FrozenObject or LearningRule) is not really necessary to get your code integrated with Nengo. If your computation fall outside of what Nengo is generally assuming (e.g., the learning rule formulation), then it is perfectly okay to create your own custom class to do the computation you need. At its very core, all Nengo is expecting is a chain of Nengo operators (see here: https://github.com/nengo/nengo/blob/master/nengo/builder/operator.py) associated with your object.

7

Perfect! Thank you so much for explaining this in detail :slight_smile:

I have one more question: is it safe to assume that the delta generated at each step by the builtin learning rules falls somewhere in the [0,1] range? Or is this just something that needs to be controlled by setting an adequate learning rate?
I ask because I’m trying to work out if I should be normalising the raw weights given by my memristors to some range (with [0,1] being the most natural].

8

There are no limits on what values delta can take in the built-in learning rules. It’s up to the learning rule to handle whatever value of delta it is given. I believe that none of the built-in learning rules enforce a limit on delta though.

However, for your learning rule, you can enforce a limit on delta if you want. But that decision is solely dependent on the desired behaviour of your learning rule.

9

Excellent, thank you!

My next question maybe has more to do with Python that with Nengo per se, but I’ve been noticing some weird behaviour when updating Signals that I can’t really pin down.

In the make_step() function of my backend class I access the delta Signal:

def make_step( self, signals, dt, rng ):
    pre_filtered = signals[ self.pre_filtered ]
    error = signals[ self.error ]
    delta = signals[ self.delta ]
    n_neurons = pre_filtered.shape[ 0 ]
    alpha = -self.learning_rate * dt / n_neurons
    encoders = signals[ self.encoders ]
    weights = signals[ self.weights ]
    pos_memristors = signals[ self.pos_memristor ]
    neg_memristors = signals[ self.neg_memristor ]

with self.delta being the getter method:

@property
def delta( self ):
    return self.updates[ 0 ]

In the actual step_simmpes() step function called by the Simulator if I do something like this, my IDE (PyCharm) says that all is fine:

delta[delta>0] = 1

while instead something like this is flagged:

delta = 1

Is it just that my IDE does not know that the variables from the outer scope make_step() are loaded by the Buider? So is any assignment like my second one actually correct i.e., can I be sure that any direct assignment in that form is going to then reflect in the weights Signal?

10

When you do delta[delta>0] = 1 that mutates the underlying data referenced by the name delta. But when you do delta = 1 that reassigns the locally-scoped variable-name delta to point to something different. The distinction can be important, for example with the following:

delta = 0
def f():
    delta = 1
f()
print(delta)

delta = [0]
def f():
    delta[:] = [1]
f()
print(delta)

Try running these two examples and note the difference. You may wish to read up on mutability and variable scoping in Python for more information. Hope this helps.

11

Yes, it does, thank you!. I was imagining something along those lines but I was convinced that in Python everything is a reference.

How would you then suggest I selectively update the delta Signal? Is what I’m doing (NumPy indexing) already the best way? Is there any way to assign a totally new array to the Signal from inside the Learning Rule?

12

Some sort of indexing would be ideal as under-the-hood that is triggering a method (the __setitem__ method in case you are curious to look into some low-level Python) on the NumPy array which mutates the referenced data.

So something like:

delta[:] = 0

will overwrite the data with all zeros. Or similarly:

delta[:] = new_delta

will replace the data that delta points to with the contents of new_delta. Likewise you could replace the : with some other indexing scheme (e.g., delta>0) to slice into the destination array.

1 Like