How can Tensorflow be used to visualize the loss versus training using Python?

TensorFlow can be used to visualize the loss versus training using the matplotlib library and plot method to plot the data. This visualization helps monitor training progress and detect issues like overfitting.

Read More: What is TensorFlow and how Keras work with TensorFlow to create Neural Networks?

A neural network that contains at least one layer is known as a convolutional layer. We can use the Convolutional Neural Network to build learning model.

The intuition behind transfer learning for image classification is, if a model is trained on a large and general dataset, this model can be used to effectively serve as a generic model for the visual world. It would have learned the feature maps, which means the user won't have to start from scratch by training a large model on a large dataset.

TensorFlow Hub is a repository that contains pre-trained TensorFlow models. TensorFlow can be used to fine-tune learning models.

We will understand how to use models from TensorFlow Hub with tf.keras, use an image classification model from TensorFlow Hub. Once this is done, transfer learning can be performed to fine-tune a model for customized image classes.

Complete Training and Visualization Example

Here's a complete example showing how to train a model and visualize the loss ?

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

# Create sample training data
X_train = np.random.random((1000, 10))
y_train = np.random.randint(0, 2, (1000, 1))

# Build a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# Compile the model
model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

# Train the model and capture history
history = model.fit(X_train, y_train, 
                    epochs=20, 
                    batch_size=32,
                    validation_split=0.2,
                    verbose=0)

print("Training completed successfully!")

Training completed successfully!

Visualizing Loss vs Training Steps

Now we can visualize the training and validation loss over epochs ?

# Extract loss values from history
train_loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(train_loss) + 1)

# Create the plot
plt.figure(figsize=(10, 6))
plt.plot(epochs, train_loss, 'b-', label='Training Loss')
plt.plot(epochs, val_loss, 'r-', label='Validation Loss')
plt.title('Loss vs Training Steps')
plt.xlabel('Training Steps (Epochs)')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.show()

print(f"Final training loss: {train_loss[-1]:.4f}")
print(f"Final validation loss: {val_loss[-1]:.4f}")

Final training loss: 0.6234
Final validation loss: 0.6187

Using Custom Callback for Batch-Level Visualization

For more detailed monitoring, you can track loss at each batch using custom callbacks ?

class LossHistoryCallback(tf.keras.callbacks.Callback):
    def __init__(self):
        self.batch_losses = []
    
    def on_batch_end(self, batch, logs=None):
        self.batch_losses.append(logs.get('loss'))

# Create callback instance
batch_stats_callback = LossHistoryCallback()

# Train with batch-level tracking
model.fit(X_train, y_train,
          epochs=5,
          batch_size=32,
          callbacks=[batch_stats_callback],
          verbose=0)

# Visualize batch-level losses
plt.figure(figsize=(12, 6))
plt.plot(batch_stats_callback.batch_losses)
plt.title('Loss vs Training Steps (Batch Level)')
plt.xlabel('Training Steps (Batches)')
plt.ylabel('Loss')
plt.ylim([0, 2])
plt.grid(True)
plt.show()

Key Benefits of Loss Visualization

Conclusion

Visualizing loss versus training steps helps monitor model performance and detect issues early. Use history.history for epoch-level tracking or custom callbacks for batch-level monitoring. This visualization is essential for effective model training and debugging.