Accessing LMU Trained Weights

Hello again,

I was wondering how you might go about accessing the trained weights of the connections inside the LMUCell (as in Legendre Memory Units in NengoDL — NengoDL 3.5.1.dev0 docs)? I know that you can access the weights and other parameters using sim.data[conn].weights or sim.get_nengo_params, but I am unsure how to apply them in this circumstance.

Thanks in advance,
Jack

Hi @jack9117,

To get at the network parameters, you’ll need a handle to the object in question in the code that’s attempting to access it. In most of our Nengo and NengoDL examples, objects like ensembles and nodes typically return a handle to them when you create the object, e.g.:

ens = nengo.Ensemble(10, 1)

In the example above, ens is the handle to the underlying Ensemble object, and you can use this ens handle to access the parameters associated with it.

In our examples, this is not the case for some of the connection objects. We sometimes just make a connection without storing a handle to it, e.g.:

nengo.Connection(inp, ens)

However, storing a handle to said connection object is as simple as assigning it to a variable. For the connection code above, you’d change it to something like this:

conn = nengo.Connection(inp, ens)

And now, conn serves as a handle to the connection between inp and ens.

The same can be done with the LMU example you linked in your post. If you wanted to get the weights of the lmu.hout connection, you’d modify this code:

nengo.Connection(lmu.h, out, transform=nengo_dl.dists.Glorot(), synapse=None)

to something like this:

conn_out = nengo.Connection(lmu.h, out, transform=nengo_dl.dists.Glorot(), synapse=None)

Once that’s done, we can get at the connection parameters like so:

with sim:
    print(sim.data[conn_out].weights)

It get’s a little more complex when dealing with objects within Nengo networks, as in the case of the LMUCell, but it’s not too different. First, we make the observation that the LMUCell is a Python class. Being a Python class, variables can be stored in the class object by using the self. object reference. And when you want to access the stored variable from outside the class, you use the .<variable_name> syntax.

As an example, to get at the hh connection from outside the LMUCell class, first, we’d change the following line of code:

    nengo.Connection(
                self.h, self.h, transform=nengo_dl.dists.Glorot(), synapse=0
            )

to this:

    self.conn_h_h = nengo.Connection(
                self.h, self.h, transform=nengo_dl.dists.Glorot(), synapse=0
            )

which then stores the handle to the connection in a variable called conn_h_h.

Now, in our Nengo model, we see that an instance of the LMUCell class has been created with the code:

    # lmu cell
    lmu = LMUCell(
        units=212,
        order=256,
        theta=train_images.shape[1],
        input_d=train_images.shape[-1],
    )

which means that we can use lmu.conn_h_h to get the handle to the hh connection contained within the LMUCell network:

with sim:
    print(sim.data[lmu.conn_h_h].weights)