Implementing a Keras Progress Bar

What you will learn

  • How to implement Keras’s model training progress bar in PyTorch

Overview

Keras implements a progress bar that shows a nice snapshot of what is occurring during model training. This offers a more visually appealing command line report and is less burdensome than a developer having to write code to gather important metrics.

Adding a Keras progress bar only takes a few lines of code and is a nice addition to any training/test loop.

Traditional Approach

Let’s start with a simple example borrowing from the below Learn the Basics Tutorials:

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

BATCH_SIZE = 64

training_data = datasets.FashionMNIST(root="data", train=True, download=True, transform=ToTensor())

test_data = datasets.FashionMNIST(root="data", train=False, download=True, transform=ToTensor())

train_dataloader = DataLoader(training_data, batch_size=BATCH_SIZE)
test_dataloader = DataLoader(test_data, batch_size=BATCH_SIZE)


class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits


model = NeuralNetwork()
print(model)

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

Define a Traditional Training Loop

def train_loop(dataloader, model, batch_size, loss_fn, optimizer):
    size = len(dataloader.dataset)
    # Set the model to training mode
    # Important for batch normalization and dropout layers
    # Unnecessary in this situation but added for best practices
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * batch_size + len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

Define a Traditional Testing Loop

def test_loop(dataloader, model, loss_fn):
    # Set the model to evaluation mode
    # Important for batch normalization and dropout layers
    # Unnecessary in this situation but added for best practices
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    # Evaluating the model with torch.no_grad()
    # ensures that no gradients are computed during test mode
    # also serves to reduce unnecessary gradient computations
    # and memory usage for tensors with requires_grad=True
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

Train the Model

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

epochs = 10
batch_size = BATCH_SIZE
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, batch_size, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")

Epoch 1
-------------------------------
loss: 2.299914  [   64/60000]
loss: 2.298036  [ 6464/60000]
loss: 2.276126  [12864/60000]
loss: 2.267775  [19264/60000]
loss: 2.249449  [25664/60000]
loss: 2.211145  [32064/60000]
loss: 2.231255  [38464/60000]
loss: 2.191056  [44864/60000]
loss: 2.174907  [51264/60000]
loss: 2.138854  [57664/60000]
Test Error:
 Accuracy: 38.1%, Avg loss: 2.141976

Epoch 2
-------------------------------
loss: 2.158404  [   64/60000]
loss: 2.150520  [ 6464/60000]
loss: 2.085808  [12864/60000]
loss: 2.093486  [19264/60000]
loss: 2.037194  [25664/60000]
loss: 1.976014  [32064/60000]
loss: 2.012927  [38464/60000]
loss: 1.925277  [44864/60000]
loss: 1.924248  [51264/60000]
loss: 1.845655  [57664/60000]
Test Error:
 Accuracy: 49.5%, Avg loss: 1.848806

Epoch 3
-------------------------------
loss: 1.892331  [   64/60000]
loss: 1.859852  [ 6464/60000]
loss: 1.735108  [12864/60000]
loss: 1.771755  [19264/60000]
loss: 1.657763  [25664/60000]
loss: 1.619588  [32064/60000]
loss: 1.649848  [38464/60000]
loss: 1.547130  [44864/60000]
loss: 1.573486  [51264/60000]
loss: 1.469762  [57664/60000]
Test Error:
 Accuracy: 59.6%, Avg loss: 1.487445

...
...
<clipped output>
...
...

Epoch 9
-------------------------------
loss: 0.879913  [   64/60000]
loss: 0.941636  [ 6464/60000]
loss: 0.724766  [12864/60000]
loss: 0.908513  [19264/60000]
loss: 0.808189  [25664/60000]
loss: 0.799041  [32064/60000]
loss: 0.880020  [38464/60000]
loss: 0.834604  [44864/60000]
loss: 0.849155  [51264/60000]
loss: 0.818514  [57664/60000]
Test Error:
 Accuracy: 69.5%, Avg loss: 0.818927

Epoch 10
-------------------------------
loss: 0.827772  [   64/60000]
loss: 0.901987  [ 6464/60000]
loss: 0.678783  [12864/60000]
loss: 0.872766  [19264/60000]
loss: 0.777142  [25664/60000]
loss: 0.760889  [32064/60000]
loss: 0.846925  [38464/60000]
loss: 0.809381  [44864/60000]
loss: 0.816719  [51264/60000]
loss: 0.789865  [57664/60000]
Test Error:
 Accuracy: 70.6%, Avg loss: 0.788775

Done!

Setup a Keras Progress Bar

The same train loop and test loop can be performed using Keras’s ProgBar class. This can cut down on the code needed to report metrics during training and be more visual appealing.

Define a Training Loop with a Progress Bar

import os

os.environ["KERAS_BACKEND"] = "torch"
import keras


def train_loop_pbar(dataloader, model, loss_fn, optimizer, epochs):
    # Move the epoch looping into the train loop
    # so epoch and pbar reported together
    for t in range(epochs):
        print(f"Epoch {t+1}/{epochs}")

        # Set the target of the pbar to
        # the number of iterations (batches) in the dataloader
        n_batches = len(dataloader)
        pbar = keras.utils.Progbar(target=n_batches)

        model.train()
        for batch, (X, y) in enumerate(dataloader):
            # Compute prediction and loss
            pred = model(X)
            loss = loss_fn(pred, y)
            correct = (pred.argmax(1) == y).type(torch.float).sum().item()
            correct /= len(y)

            # Backpropagation
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            # Update the pbar with each batch
            pbar.update(batch, values=[("loss", loss), ("acc", correct)])

            # Finish the progress bar with a final update
            # This ensures the progress bar reaches the end
            # (i.e. the target of n_batches is met on the last batch)
            if batch + 1 == n_batches:
                pbar.update(n_batches, values=None)

Define a Testing Loop with a Progress Bar

def test_loop_pbar(dataloader, model, loss_fn):
    n_batches = len(dataloader)
    pbar = keras.utils.Progbar(target=n_batches)

    model.eval()
    with torch.no_grad():
        for batch, (X, y) in enumerate(dataloader):
            pred = model(X)
            loss = loss_fn(pred, y)
            correct = (pred.argmax(1) == y).type(torch.float).sum().item()
            correct /= len(y)

            pbar.update(batch, values=[("loss", loss), ("acc", correct)])

            if batch + 1 == n_batches:
                pbar.update(n_batches, values=None)

Train the Model with a Progress Bar

# Re-instantiate model
model = NeuralNetwork()

# Train the model with progress bar
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

epochs = 10
train_loop_pbar(train_dataloader, model, loss_fn, optimizer, epochs)

Epoch 1/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 8s 8ms/step - loss: 2.2195 - acc: 0.4295
Epoch 2/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 7s 8ms/step - loss: 1.9854 - acc: 0.5649
Epoch 3/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 7s 8ms/step - loss: 1.6332 - acc: 0.5863
Epoch 4/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 7s 8ms/step - loss: 1.3331 - acc: 0.6315
Epoch 5/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 8s 9ms/step - loss: 1.1353 - acc: 0.6518
Epoch 6/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 10ms/step - loss: 1.0048 - acc: 0.6651
Epoch 7/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 8s 9ms/step - loss: 0.9169 - acc: 0.6771
Epoch 8/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 10ms/step - loss: 0.8559 - acc: 0.6902
Epoch 9/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 10ms/step - loss: 0.8116 - acc: 0.7033
Epoch 10/10
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 9ms/step - loss: 0.7777 - acc: 0.7165

# Evaluate the model with progress bar
test_loop_pbar(test_dataloader, model, loss_fn)

157/157 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - loss: 0.7776 - acc: 0.7116

Including a Validation Dataset

The progress bar can handle validation datasets as well. Start by splitting the existing training data into a new training set and validation set using an 80%-20% split. That is, 48,000 new training and 12,000 validation datapoints.

new_train_set, val_set = torch.utils.data.random_split(
    training_data, [int(len(training_data) * 0.8), int(len(training_data) * 0.2)]
)

new_train_dataloader = DataLoader(new_train_set, batch_size=BATCH_SIZE)
val_dataloader = DataLoader(val_set, batch_size=BATCH_SIZE)

print(len(train_dataloader), len(train_dataloader.dataset))
print()
print(len(new_train_dataloader), len(new_train_dataloader.dataset))
print(len(val_dataloader), len(val_dataloader.dataset))

938 60000

750 48000
188 12000

Define a Training Loop with a Progress Bar Including Validation Data

def train_loop_pbar_with_val(train_dataloader, val_dataloader, model, loss_fn, optimizer, epochs):
    for t in range(epochs):
        print(f"Epoch {t+1}/{epochs}")

        n_batches = len(train_dataloader)
        pbar = keras.utils.Progbar(target=n_batches)

        # Train step
        model.train()
        for batch, (X, y) in enumerate(train_dataloader):
            # Compute prediction and loss
            pred = model(X)
            loss = loss_fn(pred, y)
            correct = (pred.argmax(1) == y).type(torch.float).sum().item()
            correct /= len(y)

            # Backpropagation
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            # Update the pbar with each batch
            pbar.update(batch, values=[("loss", loss), ("acc", correct)])

        # Validation step
        model.eval()
        val_loss, val_correct = 0, 0
        with torch.no_grad():
            for X, y in val_dataloader:
                val_pred = model(X)
                val_loss += loss_fn(val_pred, y).item()
                val_correct += (val_pred.argmax(1) == y).type(torch.float).sum().item()

            val_loss /= len(val_dataloader)
            val_correct /= len(val_dataloader.dataset)

            # Final update now belongs to the validation data
            pbar.update(n_batches, values=[("val_loss", val_loss), ("val_acc", val_correct)])

Train the Model with a Progress Bar Including Validation Data

# Re-instantiate model
model = NeuralNetwork()

# Train the model with progress bar
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

epochs = 10
train_loop_pbar_with_val(new_train_dataloader, val_dataloader, model, loss_fn, optimizer, epochs)

Epoch 1/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 9s 11ms/step - loss: 2.2487 - acc: 0.3137 - val_loss: 2.1909 - val_acc: 0.4718
Epoch 2/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 10ms/step - loss: 2.1157 - acc: 0.5244 - val_loss: 2.0218 - val_acc: 0.5712
Epoch 3/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 10ms/step - loss: 1.8946 - acc: 0.5786 - val_loss: 1.7467 - val_acc: 0.5937
Epoch 4/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 7s 10ms/step - loss: 1.5998 - acc: 0.6092 - val_loss: 1.4567 - val_acc: 0.6229
Epoch 5/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 7s 10ms/step - loss: 1.3470 - acc: 0.6311 - val_loss: 1.2464 - val_acc: 0.6442
Epoch 6/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 10ms/step - loss: 1.1730 - acc: 0.6466 - val_loss: 1.1054 - val_acc: 0.6586
Epoch 7/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 11ms/step - loss: 1.0547 - acc: 0.6574 - val_loss: 1.0077 - val_acc: 0.6715
Epoch 8/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 11ms/step - loss: 0.9714 - acc: 0.6674 - val_loss: 0.9376 - val_acc: 0.6828
Epoch 9/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 11ms/step - loss: 0.9106 - acc: 0.6786 - val_loss: 0.8853 - val_acc: 0.6913
Epoch 10/10
750/750 ━━━━━━━━━━━━━━━━━━━━ 8s 10ms/step - loss: 0.8644 - acc: 0.6892 - val_loss: 0.8449 - val_acc: 0.7011

As a sniff test, run the original test_loop() with the validation dataset and verify the output is the same as the last epoch in the train_loop_pbar_with_val() loop:

test_loop(val_dataloader, model, loss_fn)

Test Error:
 Accuracy: 70.1%, Avg loss: 0.844919

Conclusion

Writing training loops and test loops can be verbose, especially when gathering important metrics to report. Utilizing Keras’s progress bar not only makes for a more visually appealing report, but can also simplify the amount of code needed for gathering and reporting metrics when training a model.

This recipe tutorial explored how to replace the traditional manual reporting in a training and test loop with an approach that uses Keras’s ProgBar class. It was shown how to employ a progress bar with not just a training and testing set, but also with a validation set to be used during model training.