![]() |
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.
Args | |
---|---|
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.
|
Returns | |
---|---|
None |