tfa.register_all

Stay organized with collections Save and categorize content based on your preferences.

Register TensorFlow Addons' objects in TensorFlow global dictionaries.

tfa.register_all(
    keras_objects: bool = True, custom_kernels: bool = True
) -> None

When loading a Keras model that has a TF Addons' function, it is needed for this function to be known by the Keras deserialization process.

There are two ways to do this, either do

tf.keras.models.load_model(
    "my_model.tf",
    custom_objects={"LAMB": tfa.image.optimizer.LAMB}
)

or you can do:

tfa.register_all()
tf.tf.keras.models.load_model("my_model.tf")

If the model contains custom ops (compiled ops) of TensorFlow Addons, and the graph is loaded with tf.saved_model.load, then custom ops need to be registered before to avoid an error of the type:

tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered
'...' in binary running on ... Make sure the Op and Kernel are
registered in the binary running in this process.

In this case, the only way to make sure that the ops are registered is to call this function:

tfa.register_all()
tf.saved_model.load("my_model.tf")

Note that you can call this function multiple times in the same process, it only has an effect the first time. Afterward, it's just a no-op.

keras_objects boolean, True by default. If True, register all Keras objects with tf.keras.utils.register_keras_serializable(package="Addons") If set to False, doesn't register any Keras objects of Addons in TensorFlow.
custom_kernels boolean, True by default. If True, loads all custom kernels of TensorFlow Addons with tf.load_op_library("path/to/so/file.so"). Loading the SO files register them automatically. If False doesn't load and register the shared objects files. Not that it might be useful to turn it off if your installation of Addons doesn't work well with custom ops.

None