How to Plot Model Loss During Training in TensorFlow

How you can step up your model training by plotting live the learning of your model.

Apr 13 · 3 min read
Image By Author (Logos by Keras and Tensorflow)

The Keras progress bars look nice if you are training 20 epochs, but no one wants an infinite scroll in their logs of 300 epochs progress bars (I find it disgusting). It makes it difficult to get a sense of the progress of training, and it’s just bad practice (at least if you’re training from a Jupyter Notebook).

The alternative is to have a simple plot, with train and test loss, that updates every epoch or every n steps. It’s an extremely simple implementation and it’s much more useful and insightful.

So let’s get into it.

How Does a Keras Callback Work

A Keras Callback is a class that has different functions that are executed at different times during training [1]:

  • When fit/evaluate/predict starts & ends
  • When each epoch starts & ends
  • When each training batch starts & ends
  • When each evaluation (test) batch starts & ends
  • When each inference (prediction) batch starts & ends

We will focus on the epoch functions, as we will update the plot at the end of each epoch. Within these functions you can do whatever you want, so you can let your imagination run wild and free. But let’s stick to this application for now.

Our Custom Callback

Each function receives the parameter logs, which is a dictionary containing for each metric name (accuracy, loss, etc…) the corresponding value for the epoch:

logs == {
'accuracy' : 0.98,
'loss': 0.1

To plot the training progress we need to store this data and update it to keep plotting in each new epoch. We will create a dictionary to store the metrics. Each key will correspond to a metric and have a list as its value.

On Train Begin

When the training starts we will initialize all the values. And for each epoch, we will update the metrics dictionary and update the plot.

On Epoch End

First, we store the new log values into our data structure:

Then, we create a graph for each metric, which will include the train and validation metrics. Here we clear the output of our previous epoch, generate a figure with subplots, and plot the graph for each metric, and check if there is an equivalent validation metric:

Running Our Callback

You can run this callback with any verbosity level of any other callback. As we implemented it, it will clear the output, and update the plot, so there is no need to remove logs. You may even keep the progress bar for even more interactivity.

To do this you just need to include the function we implemented in your callbacks list:

Then, when you call fit() you will get these beautiful graphs that update live:

Image by Author

To wrap up

You can now showcase your training live in a cleaner and more visual way. Small changes to your workflow like this have saved me a lot of time and improved overall satisfaction with my way of working.

This is just my implementation and there are many other useful things you can do with callbacks, so give it a try and create something beautiful!

This is the complete implementation:

Geek Culture

Proud to geek out.

Sign up for Geek Culture Hits

By Geek Culture

Subscribe to receive top 10 most read stories of Geek Culture — delivered straight into your inbox, once a week. Take a look.

By signing up, you will create a Medium account if you don’t already have one. Review our Privacy Policy for more information about our privacy practices.

Check your inbox
Medium sent you an email at to complete your subscription.


Written by


I am a Data Scientist. Let's get nerdy.

Geek Culture

A new tech publication by Start it up (


Written by


I am a Data Scientist. Let's get nerdy.

Geek Culture

A new tech publication by Start it up (

Medium is an open platform where 170 million readers come to find insightful and dynamic thinking. Here, expert and undiscovered voices alike dive into the heart of any topic and bring new ideas to the surface. Learn more

Follow the writers, publications, and topics that matter to you, and you’ll see them on your homepage and in your inbox. Explore

If you have a story to tell, knowledge to share, or a perspective to offer — welcome home. It’s easy and free to post your thinking on any topic. Write on Medium