Search⌘ K
AI Features
Home
Courses
Deep Learning with JAX and Flax

Track Model Training in JAX Using TensorBoard

Explore how to track model training in JAX by logging evaluation metrics to TensorBoard. Learn to record training and test losses and accuracies during each epoch, and understand how to manage logging with SummaryWriter to monitor model performance effectively.

We'll cover the following...

Logging evaluation metrics

We can log the evaluation metrics when training machine learning models with JAX. They were obtained at the training stage. At this point, we can log the metrics to TensorBoard. In the example below, we log the training and evaluation metrics.

Python 3.8
for epoch in range(1, num_epochs + 1):
    state, train_metrics = train_one_epoch(state, train_loader)
    training_loss.append(train_metrics['loss'])
    training_accuracy.append(train_metrics['accuracy'])
    print(f"Train epoch: {epoch}, loss: {train_metrics['loss']}, accuracy: {train_metrics['accuracy'] * 100}")
    
    test_metrics = evaluate_model(state, test_images, test_labels)
    testing_loss.append(test_metrics['loss'])
    testing_accuracy.append(test_metrics['accuracy'])
    writer.add_scalar('Loss/train', train_metrics['loss'], epoch)
    writer.add_scalar('Loss/test', test_metrics['loss'], epoch)
    
    writer.add_scalar('Accuracy/train', train_metrics['accuracy'], epoch)
    writer.add_scalar('Accuracy/test', test_metrics['accuracy'], epoch)
    print(f"Test epoch: {epoch}, loss: {test_metrics['loss']}, accuracy: {test_metrics['accuracy'] * 100}")
writer.flush()
...