Solving the TensorFlow Keras Model Loss Problem | by Chaim Rand | Towards Data Science

One of the main ingredients of a successful deep neural network, is the model loss function. At Mobileye, (officially known as Mobileye, an Intel Company), we spend a lot of time cultivating our loss functions, and fine-tuning them to the precise problems that we are trying to solve. While we, naturally, desire as much flexibility as possible when it comes to defining loss functions, it should come as no surprise that high level training frameworks and APIs, might impose certain restrictions. In this post, I will describe the challenge of defining a non-trivial model loss function when using the, high-level, TensorFlow keras model.fit() training API.
Be sure to check out some of my other posts related to TensorFlow development, covering topics such as performance profiling, debugging, and monitoring the learning process.
Given that TensorFlow has been apt to frequent updates and changes, it is important to state that this post is based on TensorFlow version 2.3.
I often hear veteran deep learning engineers, among them veteran TensorFlow developers, whine about the TensorFlow high level APIs, including the TensorFlow estimator and tf.keras.model modules. They yearn for the glorious days of yore, when one had line-level control of the training loop and could “understand what was going on” with their model. They grimace at the new APIs, and brand their users weak, and feeble-minded.
There are, no doubt, advantages to writing your own training loop: greater flexibility in building your model, much more room for being creative, and, perhaps, a deeper understanding of what’s going on. In addition, in version 2, TensorFlow has done a great job at appealing to the community of custom train loop developers, both in the quality of documentation and the API offering. But, it is important to fully understand what you are giving up on when you choose to implement a custom tensor loop.
We will focus our discussion on the tf.keras model.fit() high level API. (In version 2 of TensorFlow, there is clear favoritism towards the model.fit() API over the estimator APIs.) Here are some of the advantages to training with model.fit() over implementing a custom training loop.
Lower Entry Barrier / Ease of Development: This is perhaps the most obvious advantage. The high level APIs make it relatively easy for TensorFlow newbies to create their first training job. TensorFlow offers a wide variety of tutorials and examples, and for simple DNN projects, kicking off the training becomes a matter of “plug and play”, e.g. choosing the model, choosing the loss function, plugging in the dataset, and running model.fit(). Even for more experienced TensorFlow, who are efficient, (a.k.a. lazy), the convenience of being able to start up a training job, in just a few lines of code, cannot be underestimated.
Optimality and Correctness: This, in my view, is the primary benefit to high level APIs, in general, and, in particular, to model.fit(). When I run model.fit(), I am taking advantage of many, many hours of optimizations by TensorFlow engineers to tune the flow to its optimum. Of course, it is possible that you could, with a bit of effort, match, or maybe even beat, the runtime performance (throughput) of the high level API with a custom training loop. (As with any high level API, the need to support a wide variety of use cases, probably introduces some overhead that you might be able to cut out.) But this would require, not just the initial implementation effort, but, likely, a great deal of maintenance work to keep up with the TensorFlow improvements and optimizations introduced with each new version. Not to mention the fact that the more custom code that you include in your project, the more bug prone it becomes.
Built in Utilities for Training Management and Monitoring: There are many conveniences offered by the high level API. These include a great number of callbacks for managing your training; saving checkpoints, writing summaries, updating the learning rate, early stopping, profiling and more. Additionally, the high level API, simplifies the use of training and evaluation metrics. Metrics for monitoring the training losses are automatically defined and, you can easily request additional metrics via the model.compile() API.
Distributed Training: Using model.fit also simplifies the use of TensorFlow strategies for performing distributed training. As described in the documentation, tf.distribute.Strategy is integrated in such a way that makes it seamless to distribute training using model.fit. A custom training loop will require greater effort.
The advantages to the high level API need to be weighed against the limitations it imposes. As is often the case with regards to high level APIs, certain usages may appear to be difficult to implement, or, even impossible to implement, using model.fit. And indeed, there are some things that can only be implemented with a custom training loop. A classic example is a multi-network model such as a GAN. Other examples are if you wish to apply a custom operation during the gradient calculation (say to increase performance by decreasing bit precision, as described here), or if you wish to capture tensors for debugging purposes, as described here. In many cases, if you were to dig deep enough, you would likely find that you probably could work around the apparent limitations, though, it wouldn’t necessarily be pretty. Perhaps the solution would involve creating a custom optimizer, a custom callback, a custom metric, or a custom layer. Perhaps you would need to extend (inherit from) a TensorFlow object (e.g. tf.keras.loss, tf.keras.optimizer or tf.keras.model) and overwrite one of its methods. Perhaps you might even need to dig into the TensorFlow code and make some custom hacks (I told you it might not be pretty...).
Perhaps the most limiting constraint imposed by the high level API, and the topic of this post, pertains to how to define the model loss function. We will expand on this in the next section.
At the end of the day, the decision on our team, was to prefer, whenever possible, to adapt to the high level APIs in order to rely on the built in optimizations, and take advantage of the conveniences offered.
If you choose the route of the custom training loop, you might find this post to be useful.
Note, that in TensorFlow 2.2, an intermediate level of customization was introduced via the tf.keras.model train_step and test_step functions. This enables you to take advantage of some of the optimizations and conveniences, offered by the high level fit() routine, while also inserting some of your own customization. I highly recommend evaluating this option before choosing to go for full customization. The motivation and goodness of this API, as well as an example of how to use it, is best described in this TensorFlow guide.
The challenge of configuring the training loss in tf.keras.model fit function, is where the controversy surrounding the use of the high-level model.fit() API, reaches a boiling point. It is the ultimate example of the ways in which the high level API (seemingly) introduces restrictions on the training program. In this section, I will describe the imposed restriction and a number of different ways to overcome it. The reason I describe several, and not just one, is that none of them are perfect solutions. Each have their own limitations, some of which we will detail, and deciding which one is best for you, should depend on your specific development needs.
The standard way of configuring the loss function for training with the model.fit function, is via the model.compile function, which allows you to enter one or more (or zero) losses through the loss argument. The problem is that the loss function must have the signature loss = fn(y_true, y_pred), where y_pred is one of the outputs of the model and y_true is its corresponding label coming from the training/evaluation dataset. This is great for loss functions that are clearly dependent on a single model output tensor and a single, corresponding, label tensor. If you are lucky, not only will your model conform to this standard, but you will be able to use one of the default losses provided by tf.keras.losses. But that is not the case for just about any of our models, or loss functions. Our loss functions often depend on multiple outputs and multiple labels, and tend to be a lot more complex than the default losses offered in the API.
So what are the options? One option, of course, is to abandon the default training step and implement your own. We chose, wherever possible, to stick with the default training loop and considered the following alternatives:
The first option is to modify the outputs and labels of the model in order to conform to the required signature. If the loss must receive two tensors, y_true and y_pred, then we will flatten and concatenate all of the labels it depends on, on the one hand, and all of the outputs it depends on, on the other hand, into two corresponding tensors. This requires three changes for each loss function:
Addition of two layers to your graph, tf.keras.layers.Flatten and tf.keras.layers.Concatenate
Addition of a pre-processing routine to your dataset that combines the needed labels into a single label, with the same name as the concatenated output
Addition of a pre-processing step in the loss function to split the combined tensors back into individual tensors
This method has a few potential drawbacks that you should consider:
The extra steps will introduce some computational overhead. If your model is large, (as ours is), this overhead is negligible. But you might be sensitive to it.
If you have multiple losses and the same tensors are required for more than one loss, you will essentially be duplicating the data. Once again, if your model is large, this should be negligible, but if you are GPU memory bound, or if your training bottleneck is the training data traffic into the GPU, you might be sensitive to this.
This solution is described here, (although the examples given are somewhat trivial, and do not depend on label data). The add_loss function, essentially allows you to add any tensor you want to the loss calculation. Sounds great, no? The problem is that this loss tensor cannot rely on tensors that are outside the computation graph. Specifically, this means that any labels that the loss depends on, now need to be inserted into the graph as inputs (placeholders).
The keras documentation, includes an elegant way of handling the labels when employing the add_loss function, using an endpoint layer. Rather than invoking the add_loss on the model after it has been built, this method calls for defining a custom layer, to be placed at the end of the graph, that receives the predictions and targets as inputs, and applies the add_loss in the body of its call function. The output of the layer is the model output. Naturally, this layer needs to be removed or adjusted for running model.predict().
The steps that are required for using the add_loss option are:
Addition of input layers for each of the labels that the loss depends on
Modifying the dataset by copying or moving all relevant labels to the dictionary of features.
One drawback to consider is that this method will combine all the model losses into a single reported output loss. The default loss mechanism enables you to easily distinguish between different losses and track them separately. In particular, you can easily separate between the regularization factor of the loss and the rest of the losses. When you use add loss, you are essentially mixing all the losses together, and you will need to implement a mechanism for separating them for tracking (e.g. adding the loss tensors to the model outputs or using tf summaries on the individual loss tensors).
This option, my own personal favorite, takes the endpoint layer option one step further. Rather than calling the model.add_loss function, and outputting the model predictions, we will define the layer to actually perform the loss calculation, and output the loss result. We will define the model such that the outputs includes the calculated losses, and we will define the model losses (compile losses) to receive the outputs from the loss layers, and return their scalar values, untouched.
The advantage to this option, over the previous option, is that it enables us to easily distinguish between different losses during training, by keeping them separate. This solution requires adding a dummy loss target to the dataset for each of the model loss functions.
class LossEndPoint(Layer):    def call(self, predictions, targets):        # loss_fn is a customized loss function        loss = loss_fn(predictions, targets)        return loss    def compute_output_shape(self, input_shape):        return [1]# the compile loss, simply returns the loss scalardef compile_loss(dummy_target, y_pred):    return tf.squeeze(y_pred)
Similar to the previous solution, this method requires entering all of the labels as graph input features, and moving the labels over to the dictionary of features in the dataset. It also requires special handling for calling model.predict().
Another option, more suitable to TensorFlow 1, is to provide the loss function with all of the tensors it requires in a round about way, either by extending the tf.keras.loss class, and passing the additional tensors in the constructor, similar to what is described here (just with tensors as the parameters), or by wrapping the loss function within a context that can access all required tensors:
def get_keras_loss_fn(pred_dict, true_dict):    def keras_loss(y_true, y_pred):        loss = custom_loss_function(true_dict, pred_dict)        return loss    return keras_loss
Similar to the previous solutions, this option requires defining input layers (placeholders) for the labels, as well as moving the labels over to the dictionary of features in the dataset. Note that the loss function receives a y_true and y_pred pair, which it ignores, instead applying the loss function on the tensors that were entered to the constructor. The loss function still needs to be associated, by name, with a designated model prediction and target. You can either choose one of each, arbitrarily, or define a dummy output and label.
The advantages to this method are that it does not require adding flatten and concatenation operations, but still enables you to maintain separate losses. The one significant drawback is that this round about way of entering tensors into the graph, are disallowed in the default tf 2 execution mode, and require running in non-eager mode (i.e. calling tf.compat.v1.disable_eager_execution()), which causes some of the behaviors to default back to tf version 1.
As mentioned above, TensorFlow 2.2, introduced the option of customizing the training step of the model.fit() call by overriding the train_step function of the model class. This option is very appealing, in that it removes the requirement to conform to a specific function signature, and, essentially, avoids all the disadvantages of the options we have mentioned until now. At the same time, it enables you to take advantage of the conveniences of the tf.keras.callback utilities. But, do keep in mind what you are giving up by choosing this option, as we discussed above. For example, if you want to perform distributed training, you will need to implement the gradient sharing logic.
If you do choose this option, my strong advice would be to:
start by copying over the TensorFlow implementation of train_step, and making the minimal number of changes possible
every time you upgrade to a new TensorFlow version, refactor your custom function based on the changes that were applied to the default train_step function
Choose from the above options by carefully weighing the implications of the implementation details, advantages, and disadvantages of each one of them on your own model. If your model has many overlapping losses, the overhead of the concatenation option might be forbidding. If the ability to track the individual losses is important, you can rule out the add_loss option, and if you are risk averse, you may want to avoid the custom train step. If there is no clear right choice for you, then just flip a coin, or implement them all, (like I did).
Other than the advantages, and disadvantages we have pointed out, an additional point of comparison should be time performance. This is likely to change from model to model, and unfortunately, is hard to asses without implementing all of the solutions. For our models, (which tend to be large), the runtime of each of the options were pretty similar.
My intention in this post, as in previous posts, has been to share some of the challenges we have faced using TensorFlow, and how we overcame them. Naturally, the decisions that we came to for our projects, are not necessarily the right decisions for you. The different considerations I listed, are unlikely to be all inclusive. But I hope that sharing our considerations, and our solutions, will help you navigate your own way to success. Good luck to you all.