How to train the model with for loop instead of the built-in epochs

Question

i want to use a for loop for epochs instead of the built-in ones. Does these parts are similar.

1)

epoch=0;
while epoch in range(0, 30):
        print('Epoch', epoch)
        results = model.fit(X_train, Y_train,  batch_size=16, epochs=1, callbacks=callbacks)
        epoch +=1

         

results = model.fit(X_train, Y_train,  batch_size=16, epochs=30,  callbacks=callbacks)

Answer 1

You can do it creating a custom training function. I have created a whole set of TensorFlow 2 tutorials about it. It's simpler than it looks like.

This is the code of some generic training function:

import tensorflow as tf

# This loss and optimizer are just examples, use the one you need
loss = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.Adam(learning_rate)

@tf.function
def train_on_batch(X, Y):
    with tf.GradientTape() as tape:
        current_loss = loss(model(X), Y)
    gradients = tape.gradient(current_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return current_loss

tf.GradientTape() and tape.gradient() record the gradient calculated on the loss function. Then, optimizer.apply_gradients() produces the actual weight update (the training).

Please keep attention at the @tf.function decorator. This is a feature available in TensorFlow 2.x called Autograph. It's a super powerful trick that transforms the whole training function into a TensorFlow op, making it an order of magnitude faster. To make it work, make sure it receives numpy arrays, and that it contains exclusively TensorFlow operations.

Custom training function are extremely useful when you need to implement some non-standard training. For example, I recently trained a GAN for the imputation of missing data, and using custom training made my code easier to implement and control in all its details.

Good luck!