Migrating Keras 2 code to multi-backend Keras 3

Author: Divyashree Sreepathihalli Date created: 2023/10/23 Last modified: 2023/10/30 Description: Instructions & troubleshooting for migrating your Keras 2 code to multi-backend Keras 3.
View in Colab • GitHub source
This guide will help you migrate TensorFlow-only Keras 2 code to multi-backend Keras 3 code. The overhead for the migration is minimal. Once you have migrated, you can run Keras workflows on top of either JAX, TensorFlow, or PyTorch.
This guide has two parts:
First, lets install keras-nightly.
This example uses the TensorFlow backend (os.environ["KERAS_BACKEND"] = "tensorflow"). After you've migrated your code, you can change the "tensorflow" string to "jax" or "torch" and click "Restart runtime" in Colab, and your code will run on the JAX or PyTorch backend.
First, replace your imports:
Next, start running your tests. Most of the time, your code will execute on Keras 3 just fine. All issues you might encounter are detailed below, with their fixes.
The default value of the jit_compile argument to the Model constructor has been set to True on GPU in Keras 3. This means that models will be compiled with Just-In-Time (JIT) compilation by default on GPU.
JIT compilation can improve the performance of some models. However, it may not work with all TensorFlow operations. If you are using a custom model or layer and you see an XLA-related error, you may need to set the jit_compile argument to False. Here is a list of known issues encountered when using XLA with TensorFlow. In addition to these issues, there are some ops that are not supported by XLA.
The error message you could encounter would be as follows:
For example, the following snippet of code will reproduce the above error:
How to fix it: set jit_compile=False in model.compile(..., jit_compile=False), or set the jit_compile attribute to False, like this:
Saving to the TF SavedModel format via model.save() is no longer supported in Keras 3.
The error message you could encounter would be as follows:
The following snippet of code will reproduce the above error:
How to fix it: use model.export(filepath) instead of model.save(filepath)
Loading a TF SavedModel file via keras.models.load_model() is no longer supported If you try to use keras.models.load_model() with a TF SavedModel, you will get the following error:
The following snippet of code will reproduce the above error:
How to fix it: Use keras.layers.TFSMLayer(filepath, call_endpoint="serving_default") to reload a TF SavedModel as a Keras layer. This is not limited to SavedModels that originate from Keras – it will work with any SavedModel, e.g. TF-Hub models.
Model() can no longer be passed deeply nested inputs/outputs (nested more than 1 level deep, e.g. lists of lists of tensors).
You would encounter errors as follows:
The following snippet of code will reproduce the above error:
How to fix it: replace nested input with either dicts, lists, and tuples of input tensors.
In Keras 2, TF autograph is enabled by default on the call() method of custom layers. In Keras 3, it is not. This means you may have to use cond ops if you're using control flow, or alternatively you can decorate your call() method with @tf.function.
You would encounter an error as follows:
The following snippet of code will reproduce the above error:
How to fix it: decorate your call() method with @tf.function
Using a TF op on a Keras tensor during functional model construction is disallowed: "A KerasTensor cannot be used as input to a TensorFlow function".
The error you would encounter would be as follows:
The following snippet of code will reproduce the error:
How to fix it: use an equivalent op from keras.ops.
The evaluate() method of a multi-output model no longer returns individual output losses separately. Instead, you should utilize the metrics argument in the compile() method to keep track of these losses.
When dealing with multiple named outputs, such as output_a and output_b, the legacy tf.keras would include _loss, _loss, and similar entries in metrics. However, in keras 3.0, these entries are not automatically added to metrics. They must be explicitly provided in the metrics list for each individual output.
The following snippet of code will reproduce the above behavior:
Setting a tf.Variable as an attribute of a Keras 3 layer or model will not automatically track the variable, unlike in Keras 2. The following snippet of code will show that the tf.Variables are not being tracked.
You will see the following warning:
How to fix it: use self.add_weight() method or opt for a keras.Variable instead. If you are currently using tf.variable, you can switch to keras.Variable.
None entries are not allowed as part of nested (e.g. list/tuples) tensor arguments in Layer.call(), nor as part of call()'s nested return values.
If the None in the argument is intentional and serves a specific purpose, ensure that the argument is optional and structure it as a separate parameter. For example, consider defining the call method with optional argument.