Save best model with train_step without model.fit()

Question

Hello I'm new in machine learning, so I'm trying to save the best model weights out of 30 epochs. Now I can only save all 30 models using this code

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
    name='train_accuracy')

transformer = Transformer(num_layer, d_model, num_heads, dff, row_size, col_size, target_vocab_size,
                          max_pos_encoding=target_vocab_size, rate=dropout_rate)

@tf.function
def train_step(img_tensor, tar):
    tar_inp = tar[:, :-1]
    tar_real = tar[:, 1:]

    dec_mask = create_masks_decoder(tar_inp)

    with tf.GradientTape() as tape:
        predictions, _ = transformer(img_tensor, tar_inp,
                                     True,
                                     dec_mask)
        loss = loss_function(tar_real, predictions)

    gradients = tape.gradient(loss, transformer.trainable_variables)
    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))

    train_loss(loss)
    train_accuracy(tar_real, predictions)

for epoch in range(30):
    start = time.time()

    train_loss.reset_states()
    train_accuracy.reset_states()

    for (batch, (img_tensor, tar)) in enumerate(dataset):
        train_step(img_tensor, tar)

        if batch % 50 == 0:
            print('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(
                epoch + 1, batch, train_loss.result(), train_accuracy.result()))

    print('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1,
                                                        train_loss.result(),
                                                        train_accuracy.result()))

    print('Time taken for 1 epoch: {} secs
'.format(time.time() - start))
    model_name = 'image_caption_transformer_' + str(epoch + 1) + '.h5'
    transformer.save_weights(model_name)

I wanted to try using ModelCheckpoint from keras but I don't know how to implement it without model.fit(), any solutions to save the best model with the code above or change the code above to use model.fit()?

user_14774959 · Answer

To keep this code you can do something like this while looping over the epochs:

best_metrics = {"epoch": 0, "accuracy": 0}
for epoch in range(30):
    start = time.time()

    train_loss.reset_states()
    train_accuracy.reset_states()

    for (batch, (img_tensor, tar)) in enumerate(dataset):
        train_step(img_tensor, tar)

        if batch % 50 == 0:
            print('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(
                epoch + 1, batch, train_loss.result(), train_accuracy.result()))

    print('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1,
                                                        train_loss.result(),
                                                        train_accuracy.result()))

    print('Time taken for 1 epoch: {} secs
'.format(time.time() - start))
    if train_accuracy.result() >= best_metrics["accuracy"]:
        print("Saving...")
        model_name = 'image_caption_transformer_' + str(epoch + 1) + '.h5'
        best_metrics = {"epoch": epoch, "accuracy": train_accuracy.result()}
        transformer.save_weights(model_name)

Basically you can keep a dictionary with the best data so far. It can help you keep track of the best accuracy (or other metrics if you have them) at a given epoch. At every epoch you compare the new value with the best one saved inside the dictionary. If the current epoch has higher metrics you update the dictionary and save the weights, otherwise you move to the next epoch.

Save best model with train_step without model.fit()

1 Answers