NengoDL MNIST tutorial: trained model weights

Hello Nengo community,

I would like to extract the weights of the trained model. I have questions related to the save_params command in the NengoDl MNIST tutorial. Does this command save the weights of the model?

For instance, the network below it gives 5 NumPy arrays in the save_params file.

with nengo.Network(seed=0) as net:
    net.config[nengo.Ensemble].max_rates = nengo.dists.Choice([100])
    net.config[nengo.Ensemble].intercepts = nengo.dists.Choice([0])
    net.config[nengo.Connection].synapse = None
    neuron_type = nengo.LIF(amplitude=0.01)
    nengo_dl.configure_settings(stateful=False)
    inp = nengo.Node(np.zeros(28 * 28))  
    x = nengo_dl.Layer(tf.keras.layers.Dense(4))(inp, shape_in=(28* 28, 1))
    x = nengo_dl.Layer(neuron_type)(x)
    out = nengo_dl.Layer(tf.keras.layers.Dense(units=4))(x)
    out_p = nengo.Probe(out, label="out_p")
    out_p_filt = nengo.Probe(out, synapse=0.1, label="out_p_filt")

When I unzip the save param file it gives a warning “headers error”. But it successfully extracts 5 numpy arrays. The shape of the arrays saved by save_params are: (3136,), (1, 4), (4,), (3136, 4) and (4,) respectively. If my understanding is correct, the first layer has the dimension of 3136 because 28x28 pixels are connected to 4 neurons of the first dense layer. Then dense layer output is (1,4) and (4,) is the output of LIF. Am I right? If yes, then how is the next dimension (3136,4)?

Thank you in advance for your answer.

Hi @Choozi, and welcome back! :smiley:

To answer your questions:

The save_params function saves the parameters of the entire model. This includes the weights, biases, and state information (if specified). You can read the documentation here, or the source code here.

That is correct. If you look at the code for the save_params function, you’ll see that all it is doing is saving the keras_model.weights object to file. The keras_model.weights object itself is a Keras layer attribute, which is documented here (look for the weights attribute). Note that since NengoDL is just using the Keras weights object, NengoDL itself doesn’t have any direct control over the ordering of the parameters that appear in the keras_model.weights object (it’s up to Keras to realize this order).

If you look at the Keras documentation, it states:

The concatenation of the lists trainable_weights and non_trainable_weights (in this order).

Thus, the second (3146,4) shaped set of weights are the non-trainable weights for the first Dense layer. In the case of the Dense layer, they are the kernel weights. You can also see this by printing out the keras_model.weights object:

with nengo_dl.Simulator(net) as sim:
    print(sim.keras_model.weights)

which outputs:

[<tf.Variable 'TensorGraph/base_params/trainable_float32_3136:0' shape=(3136,) dtype=float32, numpy=array([1., 1., 1., ..., 1., 1., 1.], dtype=float32)>, 
<tf.Variable 'TensorGraph/dense/kernel:0' shape=(1, 4) dtype=float32, numpy=
array([[-0.73366153,  0.8796015 ,  0.28695   , -0.14340228]],
      dtype=float32)>, 
<tf.Variable 'TensorGraph/dense/bias:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>, 
<tf.Variable 'TensorGraph/dense_1/kernel:0' shape=(3136, 4) dtype=float32, numpy=
array([[ 0.00088362, -0.0049368 , -0.00799659,  0.04305665],
       [ 0.01649414, -0.01347676, -0.00558941,  0.00883536],
       [-0.00379217,  0.02209238, -0.02727717,  0.00426263],
       ...,
       [-0.04357701, -0.00389264,  0.03157449,  0.00664571],
       [ 0.0045911 , -0.0011935 , -0.00038958, -0.02644506],
       [-0.04076268, -0.03773354,  0.04180735, -0.02416949]],
      dtype=float32)>, 
<tf.Variable 'TensorGraph/dense_1/bias:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>]

Hello @xchoo thank you very much…!

I am glad to be back with the community.

Thank you for your detailed answer and for sharing useful links. However, I might be wrong… but the model has no non-trainable parameters since I have used only dense layers. The fact that the dense_1/kernel is (3136,4) is only because of the output of the first layer. Since the input to the layer is (784,1), only 1 value goes to the first layer and hence it makes the kernel/weights of the first layer of shape (1,4). This means that the outputs of the layer will be 784x4 = 3136, which are connected to 4 neurons of the dense_1 layer making the kernel/weights of dense_1 (3136,4). So these are not non-trainable parameters. This is my understanding… please correct me if I am wrong.

Ah, that is correct, the dense layers to not have trainable weights. I misread the output of the keras_model.weights. I don’t use TensorFlow / NengoDL very often. :sweat_smile:

Looking at your model again, I think the trainable weights (of shape (3136,)) come from the connection from the first dense layer to the neuron_type layer, rather than between the input and the first dense layer. If you comment out the neuron_type layer, the trainable weights go away.

From the nengo_dl.Layer documentation, if you create a Layer with a Nengo neuron type, it will create an ensemble of neurons, where the number of neurons will be the same as the input dimensionality to that layer. This will also create a connection between the preceding layer to the neural ensemble, and those weights are trainable by default.

@xchoo it’s ok … no worries :slight_smile:

Yes, you are right this has something to do with the neuron_type. But this confuses me more maybe because I am looking at it from conventional ANN in Keras. Don’t these dense layers represent neurons themselves? I see the neuron_type layer as an activation function… just like for example Relu.

Isn’t it just a activation fucntion?

If it is a layer, then does this ‘neuron_type’ layer has their own weights as well?

If yes, then can I say that we have approx. two times training parameters for the same model build in Keras in conventional ANN?

Looking at the documentation

A function or Keras Layer that takes the value from an input (represented as a tf.Tensor ) and maps it to some output value, or a Nengo neuron type (which will be instantiated in a Nengo Ensemble and applied to the input).

Does this mean that the weights of the neuron_type are applied to inputs before passing it through the dense layer? Sorry If I am not making any sense… :grimacing:

@xchoo @Eric @zerone can anybody of you please comment on this question?

Hello @Choozi, I haven’t been following this topic, neither I am super good at it. But it seems like

the above is untrue. In a traditional TF model as well… if you don’t mention any activation straight-ahead in Dense layers, followed by a layer of ReLU (etc.) neurons, then the connection from Dense layer to the layer of ReLU is supposed to be one to one identity connection (i.e. no weights). In fact, the weights are on the input connections to the Dense layer. Let’s wait for other experts to resolve your doubt.

Hi @Choozi,

To get back to your original question:

Looking at the output of weights again, there is a slight correct to be made:
The first (3136,) set of weights belong to the nengo_dl.Layer(neuron_type)(x) layer. That code creates a layer of 3136 LIF neurons, and each neuron has a trainable weight attached to it.

This set of weights belong to the second dense layer, which is connecting the 3136 neurons in the LIF layer to the 4 neurons in the second dense layer.

To answer your further questions:

No. The nengo_dl.Layer(neuron_type)(x) layer is bunch of activation functions, and a set of trainable weights (one for each neuron).

Yes, this is correct.

This is not correct though. If created with defaults, dense layers have a ReLU activation function, and trainable weights. In NengoDL, you can print out the trainable weights with sim.keras_model.trainable_weights:

with nengo.Network(seed=0) as net:
    inp = nengo.Node(np.zeros(28 * 28))
    x = nengo_dl.Layer(tf.keras.layers.Dense(4))(inp, shape_in=(28 * 28, 1))

with nengo_dl.Simulator(net) as sim:
    sim.keras_model.summary()
    print(sim.keras_model.trainable_weights)