How to provide tensorflow/keras callbacks

[Sum02dean]

I need to extract the weights of the best model, however since the model.fit() method is not present I am unsure where to include my keras callbacks e.g. ModelCheckpoint() - https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint

Where do I provide the callbacks?

 batch_size = 250
learning_rate = 0.002
epochs = 1


checkpoint_filepath = 'checkpoints/'
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=True,
    monitor='val_accuracy',
    mode='max',
    save_best_only=True)


# Data loaders
loader_tr = DisjointLoader(data_tr, batch_size=batch_size, epochs=epochs)
loader_te = DisjointLoader(data_te, batch_size=batch_size)


# Build general model
model = GeneralGNN(data.n_labels, activation="softmax")
#optimizer = Adam(learning_rate)
boundaries = [0, int(np.floor(0.3 * epochs))]
values = [0.0002, 0.02, 0.2]
learning_rate_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries, values)
optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate_fn)
loss_fn = CategoricalCrossentropy()

@tf.function(input_signature=loader_tr.tf_signature(), experimental_relax_shapes=True)
@tf.autograph.experimental.do_not_convert
def train_step(inputs, target):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_fn(target, predictions) + sum(model.losses)        
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    acc = tf.reduce_mean(categorical_accuracy(target, predictions))
    return loss, acc


def evaluate(loader):
    output = []
    preds = []
    step = 0
    while step < loader.steps_per_epoch:
        step += 1
        inputs, target = loader.__next__()
        pred = model(inputs, training=False)
        preds.append(pred)
        outs = (
            loss_fn(target, pred),
            tf.reduce_mean(categorical_accuracy(target, pred)),
            len(target),  # Keep track of batch size
        )
  
        output.append(outs)
        if step == loader.steps_per_epoch:
            output = np.array(output)
            return np.average(output[:, :-1], 0, weights=output[:, -1]), preds


# Run training loop        
epoch = step = 0
results = []
for batch in loader_tr:
    step += 1
    loss, acc = train_step(*batch)
    results.append((loss, acc))

    if step == loader_tr.steps_per_epoch:
        step = 0
        epoch += 1
        results_te,_ = evaluate(loader_te)
        print(
            "Ep. {} - Loss: {:.3f} - Acc: {:.3f} - Test loss: {:.3f} - Test acc: {:.3f}".format(
                epoch, *np.mean(results, 0), *results_te
            )
        )
        results = []

Thanks.

[danielegrattarola]

Hi,

here is an example of how you can train a GNN with model.fit and callbacks. You should be able to adapt it to your case pretty easily, let me know if you need further assistance.

Cheers

[jadlover]

Hi,

Will model.fit, as you have used it in the above example, work in Disjoint mode as well, provided you give the graphs in the correct format?

Thanks!

[danielegrattarola]

Yeah it should be as simple as replacing the loaders with DisjointLoader and adapting the model definition to make use of the batch index i. See here for an example of model definition

[jadlover]

Thank you very much!