How to Implement a Custom Training Loop with TensorFlow 2.x

TensorFlow General Discussion
1

Hey everyone, :smiling_face_with_three_hearts:

I’m building a deep learning project with TensorFlow 2, but I need more control over training than the model.fit() function offers. Here’s what I want to do:

  1. Track and record extra training details besides accuracy.
  2. Make my own changes to the model’s settings after each batch.
  3. Set up a learning rate that adjusts based on how well the model does on validation data.

I’ve checked TensorFlow’s resources and found something about GradientTape for custom training, but I’m stuck putting it all together. Can anyone explain how to build a custom training loop from scratch with a clear example?

Here’s a simpler breakdown of what I’m aiming for:

  • Model: A basic image classifier using convolutional neural networks (think fancy filters for images).
  • Data: A common image dataset, like lots of pictures from CIFAR-10.
  • Metrics: Instead of just accuracy, I want to follow how precise and how good at remembering things (recall) the model gets after each training round (epoch).
  • Learning Rate Schedule: If the model’s validation score (how well it does on unseen data) doesn’t improve for 3 training rounds in a row, I want to slow down the learning rate (the speed the model learns).

I also check this: https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratchruby But I have not found any solution. could anyone guide me about this? Any code examples, advice, or helpful links would be amazing!

Thanks a lot for your help! :innocent:

2

It seems you are describing a Custom Callback need rather than a custom training loop.

3

Regarding learning rate schedule, you can have a look at Keras LearningRateScheduler API.

1 Like