Keras

Keras callbacks

W&B has three callbacks for Keras, available from wandb v0.13.4. For the legacy WandbCallback scroll down.

• WandbMetricsLogger : Use this callback for Experiment Tracking. It logs your training and validation metrics along with system metrics to Weights and Biases.

• WandbModelCheckpoint : Use this callback to log your model checkpoints to Weight and Biases Artifacts.

• WandbEvalCallback: This base callback logs model predictions to Weights and Biases Tables for interactive visualization.

These new callbacks:

• Adhere to Keras design philosophy.
• Reduce the cognitive load of using a single callback (WandbCallback) for everything.
• Make it easy for Keras users to modify the callback by subclassing it to support their niche use case.

Track experiments with WandbMetricsLogger

WandbMetricsLogger automatically logs Keras’ logs dictionary that callback methods such as on_epoch_end, on_batch_end etc, take as an argument.

This tracks:

• Training and validation metrics defined in model.compile.
• System (CPU/GPU/TPU) metrics.
• Learning rate (both for a fixed value or a learning rate scheduler.

import wandb
from wandb.integration.keras import WandbMetricsLogger

# Initialize a new W&B run
wandb.init(config={"bs": 12})

# Pass the WandbMetricsLogger to model.fit
model.fit(
    X_train, y_train, validation_data=(X_test, y_test), callbacks=[WandbMetricsLogger()]
)

WandbMetricsLogger reference

Checkpoint a model using WandbModelCheckpoint

Use WandbModelCheckpoint callback to save the Keras model (SavedModel format) or model weights periodically and uploads them to W&B as a wandb.Artifact for model versioning.

This callback is subclassed from tf.keras.callbacks.ModelCheckpoint ,thus the checkpointing logic is taken care of by the parent callback.

This callback saves:

• The model that has achieved best performance based on the monitor.
• The model at the end of every epoch regardless of the performance.
• The model at the end of the epoch or after a fixed number of training batches.
• Only model weights or the whole model.
• The model either in SavedModel format or in .h5 format.

Use this callback in conjunction with WandbMetricsLogger.

import wandb
from wandb.integration.keras import WandbMetricsLogger, WandbModelCheckpoint

# Initialize a new W&B run
wandb.init(config={"bs": 12})

# Pass the WandbModelCheckpoint to model.fit
model.fit(
    X_train,
    y_train,
    validation_data=(X_test, y_test),
    callbacks=[
        WandbMetricsLogger(),
        WandbModelCheckpoint("models"),
    ],
)

WandbModelCheckpoint reference

Log checkpoints after N epochs

By default (save_freq="epoch"), the callback creates a checkpoint and uploads it as an artifact after each epoch. To create a checkpoint after a specific number of batches, set save_freq to an integer. To checkpoint after N epochs, compute the cardinality of the train dataloader and pass it to save_freq:

WandbModelCheckpoint(
    filepath="models/",
    save_freq=int((trainloader.cardinality()*N).numpy())
)

Efficiently log checkpoints on a TPU architecture

While checkpointing on TPUs you might encounter UnimplementedError: File system scheme '[local]' not implemented error message. This happens because the model directory (filepath) must use a cloud storage bucket path (gs://bucket-name/...), and this bucket must be accessible from the TPU server. We can however, use the local path for checkpointing which in turn is uploaded as an Artifacts.

checkpoint_options = tf.saved_model.SaveOptions(experimental_io_device="/job:localhost")

WandbModelCheckpoint(
    filepath="models/,
    options=checkpoint_options,
)

Visualize model predictions using WandbEvalCallback

The WandbEvalCallback is an abstract base class to build Keras callbacks primarily for model prediction and, secondarily, dataset visualization.

This abstract callback is agnostic with respect to the dataset and the task. To use this, inherit from this base WandbEvalCallback callback class and implement the add_ground_truth and add_model_prediction methods.

The WandbEvalCallback is a utility class that provides methods to:

• Create data and prediction wandb.Table instances.
• Log data and prediction Tables as wandb.Artifact.
• Log the data table on_train_begin.
• log the prediction table on_epoch_end.

The following example uses WandbClfEvalCallback for an image classification task. This example callback logs the validation data (data_table) to W&B, performs inference, and logs the prediction (pred_table) to W&B at the end of every epoch.

import wandb
from wandb.integration.keras import WandbMetricsLogger, WandbEvalCallback


# Implement your model prediction visualization callback
class WandbClfEvalCallback(WandbEvalCallback):
    def __init__(
        self, validation_data, data_table_columns, pred_table_columns, num_samples=100
    ):
        super().__init__(data_table_columns, pred_table_columns)

        self.x = validation_data[0]
        self.y = validation_data[1]

    def add_ground_truth(self, logs=None):
        for idx, (image, label) in enumerate(zip(self.x, self.y)):
            self.data_table.add_data(idx, wandb.Image(image), label)

    def add_model_predictions(self, epoch, logs=None):
        preds = self.model.predict(self.x, verbose=0)
        preds = tf.argmax(preds, axis=-1)

        table_idxs = self.data_table_ref.get_index()

        for idx in table_idxs:
            pred = preds[idx]
            self.pred_table.add_data(
                epoch,
                self.data_table_ref.data[idx][0],
                self.data_table_ref.data[idx][1],
                self.data_table_ref.data[idx][2],
                pred,
            )


# ...

# Initialize a new W&B run
wandb.init(config={"hyper": "parameter"})

# Add the Callbacks to Model.fit
model.fit(
    X_train,
    y_train,
    validation_data=(X_test, y_test),
    callbacks=[
        WandbMetricsLogger(),
        WandbClfEvalCallback(
            validation_data=(X_test, y_test),
            data_table_columns=["idx", "image", "label"],
            pred_table_columns=["epoch", "idx", "image", "label", "pred"],
        ),
    ],
)

WandbEvalCallback reference

Memory footprint details

We log the data_table to W&B when the on_train_begin method is invoked. Once it’s uploaded as a W&B Artifact, we get a reference to this table which can be accessed using data_table_ref class variable. The data_table_ref is a 2D list that can be indexed like self.data_table_ref[idx][n], where idx is the row number while n is the column number. Let’s see the usage in the example below.

Customize the callback

You can override the on_train_begin or on_epoch_end methods to have more fine-grained control. If you want to log the samples after N batches, you can implement on_train_batch_end method.

WandbCallback [legacy]

Use the W&B library WandbCallback Class to automatically save all the metrics and the loss values tracked in model.fit.

import wandb
from wandb.integration.keras import WandbCallback

wandb.init(config={"hyper": "parameter"})

...  # code to set up your model in Keras

# Pass the callback to model.fit
model.fit(
    X_train, y_train, validation_data=(X_test, y_test), callbacks=[WandbCallback()]
)

You can watch the short video Get Started with Keras and Weights & Biases in Less Than a Minute.

For a more detailed video, watch Integrate Weights & Biases with Keras. You can review the Colab Jupyter Notebook.

The WandbCallback class supports a wide variety of logging configuration options: specifying a metric to monitor, tracking of weights and gradients, logging of predictions on training_data and validation_data, and more.

Check out the reference documentation for the keras.WandbCallback for full details.

The WandbCallback

• Automatically logs history data from any metrics collected by Keras: loss and anything passed into keras_model.compile().
• Sets summary metrics for the run associated with the “best” training step, as defined by the monitor and mode attributes. This defaults to the epoch with the minimum val_loss. WandbCallback by default saves the model associated with the best epoch.
• Optionally logs gradient and parameter histogram.
• Optionally saves training and validation data for wandb to visualize.

WandbCallback reference

Frequently Asked Questions

How do I use Keras multiprocessing with wandb?

When setting use_multiprocessing=True, this error may occur:

Error("You must call wandb.init() before wandb.config.batch_size")

To work around it:

1. In the Sequence class construction, add: wandb.init(group='...').
2. In main, make sure you’re using if __name__ == "__main__": and put the rest of your script logic inside it.