Checkpoints using Nengo

General Discussion
#1

Hey, I’m trying to use checkpoints to save my best validation accuracy model to run against my test set. I’m calling it as I would in Keras but I’m getting the error shown below. Does anyone know the correct way to do this in Nengo, my code is shown below too.

Error:

AssertionError: Tried to export a function which references ‘untracked’ resource Tensor(“4329:0”, shape=(), dtype=resource). TensorFlow objects (e.g. tf.Variable) captured by functions must be ‘tracked’ by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly.

Trackable Python objects referring to this tensor (from gc.get_referrers, limited to two hops):
<tf.Variable ‘TensorGraph/saved_state/float32_66_23040:0’ shape=(66, 23040) dtype=float32>

Code:

with nengo.Network(seed=0) as net:

    # set input dimensions
    inp = nengo.Node(np.zeros(elements*num_channels))

    # convolution layer with no activation
    out = nengo_dl.Layer(tf.keras.layers.Conv1D(filters=8, kernel_size=64, strides=32, activation=None))(
        input=inp, shape_in=(elements, num_channels)
    )

    # batch norm
    out = nengo_dl.Layer(tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=.1))(
        input=out, shape_in=(dim, 8)
    )

    # average pooling layer
    out = nengo_dl.Layer(tf.keras.layers.AveragePooling1D(pool_size=dim//2))(
        input=out, shape_in=(dim, 8)
    )

    out = nengo_dl.Layer(tf.keras.layers.Flatten())(
        input=out, shape_in=(2, 8)
    )
    # first fc layer
    out = nengo_dl.Layer(tf.keras.layers.Dense(units=16, activation="relu"))(
        input=out
    )

    # final fc layer
    out = nengo_dl.Layer(tf.keras.layers.Dense(units=2, activation="softmax"))(
        input=out
    )

    # probes user for grabbing the outputs
    out_probe = nengo.Probe(out, label="out_probe")

sim = nengo_dl.Simulator(net, minibatch_size=66)

sim.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=lr),
    loss={out_probe: tf.keras.losses.SparseCategoricalCrossentropy()},
    metrics={out_probe: ["accuracy"]}
)

checkpoint = tf.keras.callbacks.ModelCheckpoint(
    filepath="model/cp.ckpt",
    monitor="val_out_probe_accuracy",
    save_best_only=True,
    mode="max",
)

sim.fit(x={inp: train}, y={out_probe: train_labels},
        epochs=epochs,
        validation_data=({inp: valid}, {out_probe: valid_labels}),
        callbacks=[checkpoint],
)

sim.close()
#2

Hi @Ricko,

I checked with the NengoDL developers and they inform me that using the tf.keras.callbacks.ModelCheckpoint function is a little iffy with NengoDL (some, but not all NengoDL features are interpretable by the ModelCheckpoint function). They suggest a couple of approaches.

The first thing to try is to use the save_weights_only=True options when calling ModelCheckpoint. Using that option only saves the weights in the model, instead of trying to serialize the entire model (which in turn tries to serialize the NengoDL objects as well). With this, the ModelCheckpoint function should work with NengoDL’s sim.fit function, although you’ll need to test this for yourself.

The other option the developers suggested (assuming the first option doesn’t work) is to write your own callback function and use NengoDL’s save functions in place of TF’s ModelCheckpoint function. NengoDL has a built-in save_params function and a corresponding load_params function that can be used to save the entirety of the NengoDL model. You can see an example of their use in this example.

Let me know if either of these two approaches work for you and also if they don’t, so I can feed back to the NengoDL developers. :smiley: