Curated Custom Callbacks
By: Pascal Janetzky
Tags: TensorFlow, Keras, Python, Callback, Deep Learning

Curated Custom Callbacks

Photo by Quino Al on Unsplash

The following callbacks are mainly input-type agnostic, meaning that they can be used for time-series, images, and text as well.

Notes

The __slots__ attribute reduces memory usage. If you modify any callback, just add any new parameter to this list.

A TensorFlow callback has access to the model it is applied to, via the self.model attribute. I use this convenient access extensively throughout my code.

Some callbacks are designed to work with TensorBoard. If you do not need this functionality, then modify — only small changes are required, mostly you will need to delete code.

ClassificationReportCallback

This short callback comes in handy when you want to regulary check on summary statistics for your training. It requires the sklearn package to generate the actual report. The arguments are a FileWriter (to write data to a TensorBoard directory), the batch size to take during prediction, and most importantly two numpy arrays, x and y.

These arrays contain sample data (x) and corresponding labels (y), based on which the classification report is created by sklearn. The class_names and labels we pass are used to filter down which classes we want to consider (e.g. maybe only the first three or last six) and to name them. The text is written to our TensorBoard directory by the file_writer we passed, every k epochs:

LayerOutputCallback

This callback lets you take the output of any intermediate network layer, write the activations to .tsv files, and the visualize them using TensorBoard or projector.tensorflow.org. The following procedure describes using it with a dense layer.

You pass sample data (x) and labels (y), the name of the layer you want to extract from (set a layer’s name when creating the model) and a (TensorBoard-)directory to write the data to.

The frequency argument states how often we write to file; the batch size is used during output extraction. We need the name argument when we write our output to file: when we e.g. extract the output of multiple layers or multiple data samples, we want to differentiate between the generated files.

The last flag controls if you want to sync your data to W&B, which is not mandatory and out of the scope for this article.

Every frequency epochs, we create a temporary model, consisting of the normal network input and our selected layer as an output. (You can think of this as cutting the network after a chosen layer). We then let this modified network predict on our sample data.

For each single sample x, we get a prediction vector x_hat. We iterate over its content, each entry is cast to string, and written to a .tsv file (vectors.tsv). In a separate file (metadata.tsv), we store any metadata, which is usually only the class label of a particular vector. This comes in handy during visualization, when we color the data points based on their labels.

If you use or modifiy this callback, I recommend checking how it works on different layers, e.g. 2D, 3D, 4D convolution, RNN, and the like. I mainly use it for dense layers deep within my networks — other types are supported, but might require slight adaption.

ConfusionMatrixCallback

This callback is built from code available at TensorFlow and writes the data to TensorBoard.

If you only need the confusion matrix as a .png, then work directly on the figure returned by plot_confusion_matrix(). Otherwise, follow along

As before, we need some sample data x and labels y. We need a FileWriter to write our data for visualization to TensorBoard, and the class names. The last parameter is used when the matrix is generated later on, it determines the labels on the plot.

The next parameters are the image name we want to save our plots as, the batch size to use during prediction, and a frequency, specifying the number of epochs to pass until we create the next plot.

The underlying matrix is generated and visualized with sklearn. These figure object could be used to store the data directly onto file (e.g. cm_epoch_10.png, with savefig()). Modify accordingly if this is your case.

If you are interested in writing this data to TensorBoard, we have to add another step: creating a TF image object. This is done in plot_to_image(), where we temporarily save our matplotlib figure to a bytes buffer, and then use the same buffer to decode it into a tensor.

The returned tensor is then written to TensorBoard with a FileWriter. Lastly, we do some clean-up:

CustomValidationMetricCallback

This callback is useful when we have a non-standard metric that we want to run on your validation data. Consider the case where we want a metric over the whole dataset, and not calculated on a per-batch basis.

In this case, we pass the validation data to the callback. At the end of an epoch the callback will be called, and our custom metric calculation code executed. In the example, we calculate the f1 score using the sklearn library.

We then log the values to the logs dict that gets passed to on_epoch_end(). This is mandatory if we have further callbacks that depend on these values (e.g. EarlyStopping). We ensure that the values are logged first by simply placing the CustomValidationMetricCallback before any dependent ones.

The second, short callback is a quick sanity check to ensure that the desired metrics are actually logged to the logs dict: