The TensorFlow Keras Summary Capture Layer | by Chaim Rand | Towards Data Science

In previous posts, I have told you about how my team at Mobileye, (officially known as Mobileye, an Intel Company), has tackled some of the challenges that came up, while using TensorFlow to train deep neural networks. In particular, I have covered topics such as performance profiling, and debugging. This post addresses an additional component of training machine learning models, that of monitoring the learning process.
Monitoring the learning process is an important, and often time-consuming, part of DNN training, during which we track a variety of tensors, metrics and statistics, in order to understand how our training is progressing, and figure out what improvements to make to model architecture and/or hyperparameters.
One of the most popular ways of doing this when training in TensorFlow, is using TensorBoard. TensorBoard supports a variety of methods for tracking and debugging the training process. On our team, we rely heavily on TensorBoard. We track losses, generate gradient histograms, and measure activation outputs. We log metrics, display confusion matrices, and generate visual images from the output data. Sometimes we use it debug intermediate operations performed by the loss function, or to measure the distribution of weights on a specific layer in the graph. We use TensorBoard extensively.
The way to use TensorBoard is by inserting “summary” commands into the training code using the tf.summary module. In version 2, TensorFlow made significant revisions to the training flow, and, in particular, to the tf summary mechanism. While this revision simplified some of the straightforward usages of tf.summary, some of the more advanced usages are not immediately obvious. For example, it is not at all clear (to me, at least), how TensorFlow intended for me to be able to log graph (non-eager) tensors, such as the distribution of the output of a specific activation layer within my graph, or the value of an arbitrary tensor within my loss function. (If you know, please don’t hesitate to share :).)
My intention, in this post, is to propose a solution for capturing summaries from graph tensors, that relies on the creation of a custom tf.keras.Layer. Along the way, I will provide a bit more insight into the problem, as well as some of the alternative solutions we considered. If you are eager to get straight to the custom summary capture layer, feel free to skip ahead.
As TensorFlow and TensorBoard continue to undergo changes and improvements, it is important to note that my story is based on trials that I ran using TensorFlow version 2.3.0.
The following brief history is partial, in the sense that it is incomplete, and that it is based on my own experience.
In version 1 of TensorFlow, we would freely scatter tf.summary ops on arbitrary tensors across our code, and would use the tf session.run() function to record the associated tensor values to event files, at predetermined intervals. As our use of TensorBoard evolved, we would also use it to record the results of general python calculations that we would make based on the graph output, such as image visualizations of model predictions overlaid on an input frame, or visualizations of the confusion matrix. Since summaries needed to be graph ops, we would do this by creating small auxiliary graphs that would define a summary op on a numpy array (such as an image). Later, we adopted the high level tf.estimator APIs, which included controls for setting the summary interval and handled the session.run calls on the summary ops for us. We also integrated the use of the estimator hooks, such as the tf.estimator.SummarySaverHook and other custom hooks, to support our wide variety of usages. Life was grand.
Along came tf.keras with its own unique blueprint for the training flow, and its own high level training API tf.keras.model.fit(). The model.fit() function did not include support for managing tf.summarys, requiring a bit more manual labor. But we quickly worked out how to enable all of our TensorBoard usages by creating custom tf.keras.callbacks, and directly configuring the execution function to collect the desired tensors, using all kinds of voodoo involving model._fit_function.fetches and model._fit_function.fetch_callbacks. No, it wasn’t pretty, but it worked. The tf.keras API also introduced a TensorBoard callback, which supported capturing a variety of popular training monitoring metrics.
Then came TensorFlow 2. Once again, the blueprint for the training flow changed, the use of the model._fit_function became obsolete, and the mechanism for recording tensors was revised. As described by TensorFlow, the new mechanism does have its advantages. Previously, recording tensors was a two step process. Now, a call to tf.summary would result in immediate recording of the tensor to the event file. In particular, this meant that the tf summary ops were no longer built into the TensorFlow computation graph. In addition, the tf.summary call could only be called on eager tensors (or raw numpy arrays). More on this in the next section.
While a full understanding of the changes that were introduced in TensorFlow 2 is beyond the scope of this post, I would like to touch on a few of the details that pertain to the topic of collecting summaries. To get the full picture, I refer you to the TensorFlow documentation.
While in tf1, the default execution mode is graph mode, a symbolic execution mode in which users define an abstract syntax tree and execute it using the TensorFlow session.run() call, TensorFlow v2 applications run eagerly. Tensors that are created within the eager execution scope, are called eager tensors, and can be accessed freely. To improve your runtime performance you can configure functions to run in graph mode by applying the tf.function qualifier to them. This is the recommended (most performant) way to configure your training step when performing training, and is what happens automatically if you rely on the high level model.fit() training API. This means that, unless you program your training step to run in (the slower) eager execution mode, all of the tensors that you defined as part of your model, will be graph (non-eager) tensors. This includes intermediate layer inputs and outputs, and intermediate tensors in your loss function. The exceptions are tensors that are defined as outputs of the tf.function, which will be eager tensors. As mentioned above, the tf2 tf.summary ops will not work on graph tensors. The TensorFlow recommendation is to debug in eager execution mode, and to perform training in graph mode. Since the tf.summaries are required for monitoring tensors, specifically during the training process, we need a solution for extracting the values of these graph tensors.
In tf1, the underlying mechanism for extracting the values of graph tensors was the tf.session object. We would provide the session.run() function with a list of graph operations and input values, and receive as output the values of the corresponding tensors. In tf2, the sessions have been replaced by functions, specifically tf.funcions when running in graph mode. Here is the change as demonstrated in code by the TensorFlow documentation.
# TensorFlow 1.Xoutputs = session.run(f(placeholder),                       feed_dict={placeholder: input})# TensorFlow 2.0outputs = f(input)
One useful property of the session.run() method was that we had quite a bit of freedom in determining the list of input ops, and thus the list of collecting tensor values. This was particularly useful for extracting summaries. At every, predetermined summary step, we could simply expand the list of input ops to include the tensors of interest. The way to do this when using tf.functions, in tf2, is not as straightforward. Since the output of a tf.function is fixed, we would need to be able to define multiple tf.functions with varying output definitions, and toggle between the tf.functions based on the training step. Or, we could just use a single tf.function that outputs the superset of all tensors of interest. But this might come at a considerable performance cost.
The main goal of this post is to propose a solution for extracting graph tensors that is based on defining a custom tf.keras layer. But before we get into that, let’s review a couple of alternative solutions.
One, relatively easy, way to solve this is to simply disable the eager execution by calling tf.compat.v1.disable_eager_execution at the beginning of your script. As stated in the documentation, this function is intended for projects for which the migration from tf1 to tf2 proves difficult. When this function is used, the training loop essentially falls back to the legacy, tf1, training loop. In particular, this means that all the voodoo we described above for extracting graph tensors in tf.keras v1, can be carried over to tf2.
The feeling on our team was that we did not go to all the trouble of updating our code to tf2, only to fall back to using a legacy, tf1 version, training sequence. Our working assumption is that newer is better when it comes to TensorFlow versions, and that staying up do date with the latest tf version, will ensure that we are benefiting from the most up to date enhancements and optimizations. Not to mention that there are a number of features, such as tf profiling and custom training steps, that won’t work in the legacy mode.
I think it goes without saying that this option is ill-advised.
TensorFlow provides support for running a fully customized training loop (instead of model.fit()), as well as APIs for customizing just the train_step of model.fit() as described here. Here we refer to the second mode of customization, although some of the conclusions are relevant to the first as well.
It is very tempting to try to take advantage of the relative freedom that customization provides in order to support the capturing of graph tensors. The thing is, that even when using a custom training step, the advised way to do so, is such that it is wrapped by a tf.function qualifier. Of course, this makes all of the tensors part of the tf.Graph and thus inaccessible.
However, using a custom training step, does grant us some freedom in defining the tensors that are output from the graph, and thus accessible. For example, we can extend the list of model outputs to include outputs of layers that we wish to record. Or, we can define our loss function such that in addition to calculating and returning the loss, it also returns tensors that we would like to record, either in the form of images or histograms.
In the simplified code block below, we define a customized training step, such that it collects the model outputs, and the tensors returned by the loss function, and stores references to them in the extended Model class. These tensors are eager tensors, that can be freely referenced in other locations of the code. For example, we could define a custom keras callback that, at fixed intervals, writes these tensors to an event file, as histograms or images.
class CustomModel(keras.Model):    @tf.function    def train_step(self, data):        x, y = data        with tf.GradientTape() as tape:            y_pred = self(x, training=True)  # Forward pass            loss, histograms, images = self.compiled_loss(y, y_pred)        # Compute gradients        train_vars = self.trainable_variables        gradients = tape.gradient(loss, train_vars)        # Update weights        self.optimizer.apply_gradients(zip(gradients, train_vars))        return y_pred, loss, histograms, images    def make_train_function(self):        if self.train_function is not None:            return self.train_function        def train_function(iterator):            data = next(iterator)            # store references to the model outputs, loss images            # and loss histograms so that these can be recorded            # to TensorBoard            self.outputs, loss, self.histograms, self.images  =                                              self.train_step(data)            return {'loss': loss}        self.train_function = train_function        return self.train_function
You may have picked up on the built-in inefficiency of this proposal. In most cases, we want to post summaries at specific training step intervals, certainly not at every step. However, we have redefined our model such that the training step outputs the tensors of interest at every step. If we were to define the set of prediction tensors that we are interested in at every step as A, and the superset that includes the summary tensors as B, then we have expanded our tf.function to output B. This can create a significant overhead, in particular, if we have many tensors of interest. Ideally, we would want to output these tensors only at the training steps at which they will be recorded. Ideally, we would want to define two tf.functions, one that outputs A and one that outputs B, and use them interchangeably. However, as of the time of this writing, I have not discovered an elegant way to do this.
The solution we propose is based on the observations that layer weights, including non-trainable weights, as described here, are eager tensors. Thus, in order to store the value of a given graph tensor, all we need to do is assign it to a “non-trainable” weight. Let’s start by demonstrating how this can be done in order to record the value of the output of an activation layer. In the block of code below, we have created a custom tf.keras layer that extends the standard tf.keras Activation with a layer that includes a non-trainable weight, called record_tensor. The call function is enhanced, such that on every step the record_tensor field is updated with the value of the current activation result. Since it is an eager tensor, record_tensor can be read outside of the training loop, and, in particular, recorded to TensorBoard.
from tensorflow.keras.layers import Activationclass ActivationRecorder(Activation):    def __init__(self, shape, dtype, **kwargs):        self.shape = shape        self.record_tensor = tf.Variable(            shape=shape,            # initialize with batch size 1 since batch_size is                  # unknown, and set validate_shape=False                       initial_value=tf.zeros(shape=[1]+shape[1:],                                   dtype=dtype),            validate_shape=False,            dtype=dtype,            trainable=False)        super(ActivationRecorder, self).__init__(dtype=dtype,                                                  **kwargs)def capture(self,activation):        self.record_tensor.assign(activation)def call(self, inputs):        output = super(ActivationRecorder, self).call(inputs)        self.capture(output)        return outputdef get_config(self):        config = super(ActivationRecorder, self).get_config()        config.update({"shape": self.shape})        return config# when building the graph maintain a reference to the# ActivationRecorder layersactivation_layer = ActivationRecorder(shape=activation_input.shape,                                      dtype=activation_input.dtype,                                      activation='relu')activation_output = activation_layer(activation_input)... #build rest graph# train the model (model.fit())# access recorded activation as needed:# Create a file writer for the log directory.file_writer = tf.summary.create_file_writer(logdir)with file_writer.as_default():    tf.summary.histogram("activation",                          activation_layer.record_tensor,                          step=step)
We have demonstrated the tensor capturing method on an activation layer, but we could have just as easily extended just about any keras layer and added a non-trainable tensor capturing variable. (One exception appears to be the keras InputLayer. It seems that this layer is handled differently, and that the “call” function is not called for every instance and every version of TensorFlow. In order to capture graph inputs, consider using a dedicated capturing layer, as we describe below.)
An alternative solution, is to define a general purpose tensor capturing layer as described in the next section.
To support the general case of capturing a graph tensor, we can define a general purpose tensor capture layer. This layer is defined as a pass through for the input, where all it does is store the current value to its internal non-trainable weight variable. This can be used to capture tensors at any stage in the graph, including layer inputs and outputs, and tensors in the loss function. Additionally, it can be used to capture tensors in the input pipeline, which was not covered by the previous solutions.
class SummaryCaptureLayer(Layer):    def __init__(self, shape, dtype, **kwargs):        self.shape = shape        self.record_tensor = tf.Variable(            shape=shape,            # initialize with batch size 1 since batch_size             # is unknown, nnd set validate_shape=False            initial_value=tf.zeros(shape=[1]+shape[1:],                                   dtype=dtype),            validate_shape=False,            dtype=dtype,            trainable=False)        super(SummaryCaptureLayer, self).__init__(dtype=dtype,                                                   **kwargs)    def capture(self,inputs):        self.record_tensor.assign(inputs)    def call(self, inputs, **kwargs):        self.capture(inputs)        return inputs    def get_config(self):        config = super(SummaryCaptureLayer, self).get_config()        config.update({"shape": self.shape})        return config
The complete solution requires storing a reference to the created layer and then accessing the record_tensor field, as needed, as shown in the previous example.
One detail that you may have noticed, in both examples, is that in the definition of the non-trainable tf.Variable, we used an arbitrary batch_size (of 1) in the initialization value, and set validate_shape to False. This is because, while the initial value requires a well defined shape (it can’t include None in any of the dimensions), at the time of model creation, we don’t always know what the batch_size will be. (For non-input layers, this also enables the use of JSON model configuration without fixing the batch size.)
In this post we have chosen to call our custom layer a SummaryCaptureLayer since the usage we have demonstrated has been recording tensor summaries to TensorBoard. But it would have been more appropriate to call the custom layer by the name TensorCaptureLayer, since all we are actually doing is capturing the tensor. Capturing tensors can support additional needs such as debugging.
In a previous post, I expanded on some tools and techniques for debugging a broken training session. The custom layer that we described can be used as an additional debugging tool. Suppose you suspect that the root cause of your bug is a problem with one, or more, particular tensors, either in the input pipeline, the loss function, or one of your layers. By feeding the suspect tensors into the Custom Summary Layer, you can track the value of the tensors at each iteration, or more realistically, dump the tensor values on the iteration that causes the training process to fail. It is clear to see how, compared to other debugging options, this kind of debugging scheme, coupled with the proposal I shared for debugging in the previous post, could save considerable time in solving the types of monster bugs I describe in the post.
Naturally, adding custom summary layers to the model, however thin they are, will incur a performance penalty. Since the performance penalty will depend directly on the model architecture, the number of summary layers that are inserted, and the frequency at which the tensors are written to the event file, it should be evaluated on a case by case basis. However, to provide some measure of comparison, we will provide the results of a simple case study, in which we capture 49 arbitrary tensors of a deep neural network, using the custom summary capture callback. The measurements below are averaged across epochs of 100 steps during which the summaries were captured, but not written (so as to remove the overhead of the summary writer).
The relatively high overhead of outputting all of the tensors of interest using a custom train step is apparent. We also see that there is a little overhead to the use of custom summary layers. While collecting tensors using the non-eager execution mode is slightly faster than using custom summary layers, I do not believe that this advantage is great enough to warrant falling back to a legacy training scheme.
Again, these results provide little information about the performance overhead of other use cases. Each model should be evaluated independently before settling on the number and types of tensors to collect, and the frequency at which they should be collected.
In this post I have proposed a way to collect summaries of graph tensors in TensorFlow 2. The solution is not perfect, but in the absence of a better one, I think it is simple and elegant. If you know of a better solution, or if TensorFlow provides a built in solution, I would love to hear about it.