Register keras objects in the function register_all, not at import time

[gabrieldemarmiesse]

Currently, it's possible to do

import tensorflow as tf
import tensorflow_addons as tfa

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

which is not readable and will throw an error with all main linting tools (import unused).
Being explicit and readable come first in python, and similarely to how we do it with custom ops, register_all() should be called before loading a model.

See background #1151

[ashutosh1919]

@gabrieldemarmiesse , I tried calling register_all() in resource_loader.py but I think it is not the file which is executed at time of loading model. Instead, it is executed when tfa is imported. Is there any code in tfa which is executed when tf loads model.

[gabrieldemarmiesse]

@ashutosh1919 , thanks for giving a hand. I don't believe it's the right approach. The right one would be to call register_keras_object in the register all and not at import time (currently done with decorators). That requires an important refactoring with multiple steps because this is linked to many pieces. Notably the biggest piece we need to do before is to make our own layer_test (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/utils/test_utils.py#L34). Currently we grab it from the TF private API which is bad. If you want to help on this issue, a first good step would be to copy the layer_test implementation to addons and remove all the code in there that we don't need (I looked at the source code recently and the code could need a bit of cleaning).

[ashutosh1919]

@gabrieldemarmiesse , I have taken a good look at layer_test code here. Can you please explain in 1-2 lines, what part of it I should focus on understanding for cleaning. And what part needs cleaning?

[gabrieldemarmiesse]

Yes, notably, we need to adapt it to addons. This function should look for bugs in addons (the layer implementation), not tensorflow.

For example, I see that layer_test runs the layer once in a sequential model and once in a functional model. We shouldn't do that. We shouldn't even need to put it in a model to check that it works correctly. We can do my_layer(...) and check that this works. If then it breaks in a model, that's not a bug in addons, it's a bug in tensorflow. Same as testing get_weights and set_weights. That's a method in tensorflow, not in addons, so we don't need to verify it.

[seanpmorgan]

TensorFlow Addons is transitioning to a minimal maintenance and release mode. New features will not be added to this repository. For more information, please see our public messaging on this decision:
TensorFlow Addons Wind Down

Please consider sending feature requests / contributions to other repositories in the TF community with a similar charters to TFA:
Keras
Keras-CV
Keras-NLP