From 3de5fcaca9b65161ca0bcd5a0242819d8754e591 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Wed, 18 May 2022 17:39:00 -0700 Subject: [PATCH 01/14] Point optimizer to tf.keras.optimizer.legacy.Optimizer to be compatible with Keras optimizer migration --- tensorflow_addons/optimizers/adabelief.py | 2 +- tensorflow_addons/optimizers/average_wrapper.py | 11 ++++++++--- tensorflow_addons/optimizers/cocob.py | 2 +- tensorflow_addons/optimizers/conditional_gradient.py | 2 +- .../optimizers/cyclical_learning_rate.py | 8 ++++---- .../optimizers/discriminative_layer_training.py | 6 +++--- tensorflow_addons/optimizers/lamb.py | 2 +- tensorflow_addons/optimizers/lookahead.py | 9 ++++++--- tensorflow_addons/optimizers/moving_average.py | 2 +- tensorflow_addons/optimizers/novograd.py | 2 +- tensorflow_addons/optimizers/proximal_adagrad.py | 2 +- tensorflow_addons/optimizers/rectified_adam.py | 2 +- tensorflow_addons/optimizers/tests/standard_test.py | 2 +- .../optimizers/tests/weight_decay_optimizers_test.py | 2 +- .../optimizers/weight_decay_optimizers.py | 6 ++++-- tensorflow_addons/optimizers/yogi.py | 2 +- tensorflow_addons/utils/test_utils.py | 1 + 17 files changed, 37 insertions(+), 26 deletions(-) diff --git a/tensorflow_addons/optimizers/adabelief.py b/tensorflow_addons/optimizers/adabelief.py index 33297f4a33..77b6b98bed 100644 --- a/tensorflow_addons/optimizers/adabelief.py +++ b/tensorflow_addons/optimizers/adabelief.py @@ -21,7 +21,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class AdaBelief(tf.keras.optimizers.Optimizer): +class AdaBelief(tf.keras.optimizers.legacy.Optimizer): """Variant of the Adam optimizer. It achieves fast convergence as Adam and generalization comparable to SGD. diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index 0ff4992441..3673b0f867 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -22,7 +22,9 @@ from typeguard import typechecked -class AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta): +class AveragedOptimizerWrapper( + tf.keras.optimizers.legacy.Optimizer, metaclass=abc.ABCMeta +): @typechecked def __init__( self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs @@ -32,9 +34,12 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if not isinstance(optimizer, tf.keras.optimizers.Optimizer): + if not isinstance( + optimizer, + (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), + ): raise TypeError( - "optimizer is not an object of tf.keras.optimizers.Optimizer" + "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer" ) self._optimizer = optimizer diff --git a/tensorflow_addons/optimizers/cocob.py b/tensorflow_addons/optimizers/cocob.py index 9af39d5d9f..e10c971924 100644 --- a/tensorflow_addons/optimizers/cocob.py +++ b/tensorflow_addons/optimizers/cocob.py @@ -19,7 +19,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class COCOB(tf.keras.optimizers.Optimizer): +class COCOB(tf.keras.optimizers.legacy.Optimizer): """Optimizer that implements COCOB Backprop Algorithm Reference: diff --git a/tensorflow_addons/optimizers/conditional_gradient.py b/tensorflow_addons/optimizers/conditional_gradient.py index 4e6592ad07..65d7ae12f4 100644 --- a/tensorflow_addons/optimizers/conditional_gradient.py +++ b/tensorflow_addons/optimizers/conditional_gradient.py @@ -22,7 +22,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class ConditionalGradient(tf.keras.optimizers.Optimizer): +class ConditionalGradient(tf.keras.optimizers.legacy.Optimizer): """Optimizer that implements the Conditional Gradient optimization. This optimizer helps handle constraints well. diff --git a/tensorflow_addons/optimizers/cyclical_learning_rate.py b/tensorflow_addons/optimizers/cyclical_learning_rate.py index 905aad0c45..f4ec176977 100644 --- a/tensorflow_addons/optimizers/cyclical_learning_rate.py +++ b/tensorflow_addons/optimizers/cyclical_learning_rate.py @@ -58,7 +58,7 @@ def __init__( ``` You can pass this schedule directly into a - `tf.keras.optimizers.Optimizer` as the learning rate. + `tf.keras.optimizers.legacy.Optimizer` as the learning rate. Args: initial_learning_rate: A scalar `float32` or `float64` `Tensor` or @@ -146,7 +146,7 @@ def __init__( ``` You can pass this schedule directly into a - `tf.keras.optimizers.Optimizer` as the learning rate. + `tf.keras.optimizers.legacy.Optimizer` as the learning rate. Args: initial_learning_rate: A scalar `float32` or `float64` `Tensor` or @@ -215,7 +215,7 @@ def __init__( ``` You can pass this schedule directly into a - `tf.keras.optimizers.Optimizer` as the learning rate. + `tf.keras.optimizers.legacy.Optimizer` as the learning rate. Args: initial_learning_rate: A scalar `float32` or `float64` `Tensor` or @@ -286,7 +286,7 @@ def __init__( ``` You can pass this schedule directly into a - `tf.keras.optimizers.Optimizer` as the learning rate. + `tf.keras.optimizers.legacy.Optimizer` as the learning rate. Args: initial_learning_rate: A scalar `float32` or `float64` `Tensor` or diff --git a/tensorflow_addons/optimizers/discriminative_layer_training.py b/tensorflow_addons/optimizers/discriminative_layer_training.py index 22fb19798d..792d3ec092 100644 --- a/tensorflow_addons/optimizers/discriminative_layer_training.py +++ b/tensorflow_addons/optimizers/discriminative_layer_training.py @@ -21,7 +21,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class MultiOptimizer(tf.keras.optimizers.Optimizer): +class MultiOptimizer(tf.keras.optimizers.legacy.Optimizer): """Multi Optimizer Wrapper for Discriminative Layer Training. Creates a wrapper around a set of instantiated optimizer layer pairs. @@ -30,7 +30,7 @@ class MultiOptimizer(tf.keras.optimizers.Optimizer): Each optimizer will optimize only the weights associated with its paired layer. This can be used to implement discriminative layer training by assigning different learning rates to each optimizer layer pair. - `(tf.keras.optimizers.Optimizer, List[tf.keras.layers.Layer])` pairs are also supported. + `(tf.keras.optimizers.legacy.Optimizer, List[tf.keras.layers.Layer])` pairs are also supported. Please note that the layers must be instantiated before instantiating the optimizer. Args: @@ -130,7 +130,7 @@ def get_config(self): @classmethod def create_optimizer_spec( cls, - optimizer: tf.keras.optimizers.Optimizer, + optimizer: tf.keras.optimizers.legacy.Optimizer, layers_or_model: Union[ tf.keras.Model, tf.keras.Sequential, diff --git a/tensorflow_addons/optimizers/lamb.py b/tensorflow_addons/optimizers/lamb.py index b166657251..64fa6894a9 100644 --- a/tensorflow_addons/optimizers/lamb.py +++ b/tensorflow_addons/optimizers/lamb.py @@ -29,7 +29,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class LAMB(tf.keras.optimizers.Optimizer): +class LAMB(tf.keras.optimizers.legacy.Optimizer): """Optimizer that implements the Layer-wise Adaptive Moments (LAMB). See paper [Large Batch Optimization for Deep Learning: Training BERT diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index 70d732979e..d77fdb81f8 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -20,7 +20,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class Lookahead(tf.keras.optimizers.Optimizer): +class Lookahead(tf.keras.optimizers.legacy.Optimizer): """This class allows to extend optimizers with the lookahead mechanism. The mechanism is proposed by Michael R. Zhang et.al in the paper @@ -71,9 +71,12 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if not isinstance(optimizer, tf.keras.optimizers.Optimizer): + if not isinstance( + optimizer, + (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), + ): raise TypeError( - "optimizer is not an object of tf.keras.optimizers.Optimizer" + "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer" ) self._optimizer = optimizer diff --git a/tensorflow_addons/optimizers/moving_average.py b/tensorflow_addons/optimizers/moving_average.py index 79432ded82..9ef010478f 100644 --- a/tensorflow_addons/optimizers/moving_average.py +++ b/tensorflow_addons/optimizers/moving_average.py @@ -55,7 +55,7 @@ def __init__( r"""Construct a new MovingAverage optimizer. Args: - optimizer: str or `tf.keras.optimizers.Optimizer` that will be + optimizer: str or `tf.keras.optimizers.legacy.Optimizer` that will be used to compute and apply gradients. average_decay: float. Decay to use to maintain the moving averages of trained variables. diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py index 159ec55f0b..bc29f31e88 100644 --- a/tensorflow_addons/optimizers/novograd.py +++ b/tensorflow_addons/optimizers/novograd.py @@ -22,7 +22,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class NovoGrad(tf.keras.optimizers.Optimizer): +class NovoGrad(tf.keras.optimizers.legacy.Optimizer): """Optimizer that implements NovoGrad. The NovoGrad Optimizer was first proposed in [Stochastic Gradient diff --git a/tensorflow_addons/optimizers/proximal_adagrad.py b/tensorflow_addons/optimizers/proximal_adagrad.py index c8e1aa0ff2..d62f0669d4 100644 --- a/tensorflow_addons/optimizers/proximal_adagrad.py +++ b/tensorflow_addons/optimizers/proximal_adagrad.py @@ -23,7 +23,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class ProximalAdagrad(tf.keras.optimizers.Optimizer): +class ProximalAdagrad(tf.keras.optimizers.legacy.Optimizer): """Optimizer that implements the Proximal Adagrad algorithm. References: diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py index c482d56d10..0c540bfdf1 100644 --- a/tensorflow_addons/optimizers/rectified_adam.py +++ b/tensorflow_addons/optimizers/rectified_adam.py @@ -21,7 +21,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class RectifiedAdam(tf.keras.optimizers.Optimizer): +class RectifiedAdam(tf.keras.optimizers.legacy.Optimizer): """Variant of the Adam optimizer whose adaptive learning rate is rectified so as to have a consistent variance. diff --git a/tensorflow_addons/optimizers/tests/standard_test.py b/tensorflow_addons/optimizers/tests/standard_test.py index f1d284ad68..948ef6d2b0 100644 --- a/tensorflow_addons/optimizers/tests/standard_test.py +++ b/tensorflow_addons/optimizers/tests/standard_test.py @@ -33,7 +33,7 @@ classes_to_test = discover_classes( - optimizers, tf.keras.optimizers.Optimizer, class_exceptions + optimizers, tf.keras.optimizers.legacy.Optimizer, class_exceptions ) diff --git a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py index edf3b97e0e..2ba5a3944c 100644 --- a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py +++ b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py @@ -406,7 +406,7 @@ def test_var_list_with_exclude_list_sgdw(dtype): [ weight_decay_optimizers.SGDW, weight_decay_optimizers.extend_with_decoupled_weight_decay( - tf.keras.optimizers.SGD + tf.keras.optimizers.legacy.SGD ), ], ) diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py index c4fbd60e5a..e0405b81ba 100644 --- a/tensorflow_addons/optimizers/weight_decay_optimizers.py +++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py @@ -263,8 +263,10 @@ def _do_use_weight_decay(self, var): @typechecked def extend_with_decoupled_weight_decay( - base_optimizer: Type[tf.keras.optimizers.Optimizer], -) -> Type[tf.keras.optimizers.Optimizer]: + base_optimizer: Type[ + Union[tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer] + ], +) -> Type[Union[tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer]]: """Factory function returning an optimizer class with decoupled weight decay. diff --git a/tensorflow_addons/optimizers/yogi.py b/tensorflow_addons/optimizers/yogi.py index 0d4389ac1c..4d602f71f4 100644 --- a/tensorflow_addons/optimizers/yogi.py +++ b/tensorflow_addons/optimizers/yogi.py @@ -50,7 +50,7 @@ def _solve(a, b, c): @tf.keras.utils.register_keras_serializable(package="Addons") -class Yogi(tf.keras.optimizers.Optimizer): +class Yogi(tf.keras.optimizers.legacy.Optimizer): """Optimizer that implements the Yogi algorithm in Keras. See Algorithm 2 of diff --git a/tensorflow_addons/utils/test_utils.py b/tensorflow_addons/utils/test_utils.py index fde88c6a26..e137c19e98 100644 --- a/tensorflow_addons/utils/test_utils.py +++ b/tensorflow_addons/utils/test_utils.py @@ -25,6 +25,7 @@ from tensorflow_addons import options from tensorflow_addons.utils import resource_loader + if tf.__version__[:3] > "2.8": from keras.testing_infra.test_utils import layer_test # noqa: F401 else: From 3bdda5e9c5f9815afa92942ad1c629a5d510541b Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Wed, 18 May 2022 17:49:10 -0700 Subject: [PATCH 02/14] small fix --- tensorflow_addons/optimizers/average_wrapper.py | 3 ++- tensorflow_addons/optimizers/lookahead.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index 3673b0f867..44fe485ccb 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -39,7 +39,8 @@ def __init__( (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), ): raise TypeError( - "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer" + "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " + "or tf.keras.optimizers.Optimizer." ) self._optimizer = optimizer diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index d77fdb81f8..cec9b57f38 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -76,7 +76,8 @@ def __init__( (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), ): raise TypeError( - "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer" + "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " + "or tf.keras.optimizers.Optimizer." ) self._optimizer = optimizer From 27a9a761e22426dbca343813082994e1b8f38548 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Mon, 23 May 2022 12:11:21 -0700 Subject: [PATCH 03/14] add version control --- tensorflow_addons/optimizers/adabelief.py | 7 ++++++- tensorflow_addons/optimizers/average_wrapper.py | 10 +++++++--- tensorflow_addons/optimizers/cocob.py | 8 +++++++- .../optimizers/discriminative_layer_training.py | 8 +++++++- tensorflow_addons/optimizers/lamb.py | 7 ++++++- tensorflow_addons/optimizers/lazy_adam.py | 8 +++++++- tensorflow_addons/optimizers/lookahead.py | 8 +++++++- tensorflow_addons/optimizers/novograd.py | 8 +++++++- tensorflow_addons/optimizers/proximal_adagrad.py | 8 +++++++- tensorflow_addons/optimizers/rectified_adam.py | 7 ++++++- tensorflow_addons/optimizers/yogi.py | 8 +++++++- 11 files changed, 74 insertions(+), 13 deletions(-) diff --git a/tensorflow_addons/optimizers/adabelief.py b/tensorflow_addons/optimizers/adabelief.py index 77b6b98bed..8c3317c4e9 100644 --- a/tensorflow_addons/optimizers/adabelief.py +++ b/tensorflow_addons/optimizers/adabelief.py @@ -19,9 +19,14 @@ from typing import Union, Callable, Dict +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + @tf.keras.utils.register_keras_serializable(package="Addons") -class AdaBelief(tf.keras.optimizers.legacy.Optimizer): +class AdaBelief(optimizer_class): """Variant of the Adam optimizer. It achieves fast convergence as Adam and generalization comparable to SGD. diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index 44fe485ccb..268138435a 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -22,9 +22,13 @@ from typeguard import typechecked -class AveragedOptimizerWrapper( - tf.keras.optimizers.legacy.Optimizer, metaclass=abc.ABCMeta -): +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + +class AveragedOptimizerWrapper(optimizer_class, metaclass=abc.ABCMeta): @typechecked def __init__( self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs diff --git a/tensorflow_addons/optimizers/cocob.py b/tensorflow_addons/optimizers/cocob.py index e10c971924..ad2c3e5178 100644 --- a/tensorflow_addons/optimizers/cocob.py +++ b/tensorflow_addons/optimizers/cocob.py @@ -18,8 +18,14 @@ import tensorflow as tf +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class COCOB(tf.keras.optimizers.legacy.Optimizer): +class COCOB(optimizer_class): """Optimizer that implements COCOB Backprop Algorithm Reference: diff --git a/tensorflow_addons/optimizers/discriminative_layer_training.py b/tensorflow_addons/optimizers/discriminative_layer_training.py index 792d3ec092..823c3fc0be 100644 --- a/tensorflow_addons/optimizers/discriminative_layer_training.py +++ b/tensorflow_addons/optimizers/discriminative_layer_training.py @@ -20,8 +20,14 @@ from typeguard import typechecked +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class MultiOptimizer(tf.keras.optimizers.legacy.Optimizer): +class MultiOptimizer(optimizer_class): """Multi Optimizer Wrapper for Discriminative Layer Training. Creates a wrapper around a set of instantiated optimizer layer pairs. diff --git a/tensorflow_addons/optimizers/lamb.py b/tensorflow_addons/optimizers/lamb.py index 64fa6894a9..d96129498a 100644 --- a/tensorflow_addons/optimizers/lamb.py +++ b/tensorflow_addons/optimizers/lamb.py @@ -27,9 +27,14 @@ from tensorflow_addons.utils.types import FloatTensorLike from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + @tf.keras.utils.register_keras_serializable(package="Addons") -class LAMB(tf.keras.optimizers.legacy.Optimizer): +class LAMB(optimizer_class): """Optimizer that implements the Layer-wise Adaptive Moments (LAMB). See paper [Large Batch Optimization for Deep Learning: Training BERT diff --git a/tensorflow_addons/optimizers/lazy_adam.py b/tensorflow_addons/optimizers/lazy_adam.py index 7b4cf8e201..aac7a38706 100644 --- a/tensorflow_addons/optimizers/lazy_adam.py +++ b/tensorflow_addons/optimizers/lazy_adam.py @@ -27,8 +27,14 @@ from typing import Union, Callable +if tf.__version__[:3] > "2.8": + adam_optimizer_class = tf.keras.optimizers.legacy.Adam +else: + adam_optimizer_class = tf.keras.optimizers.Adam + + @tf.keras.utils.register_keras_serializable(package="Addons") -class LazyAdam(tf.keras.optimizers.Adam): +class LazyAdam(adam_optimizer_class): """Variant of the Adam optimizer that handles sparse updates more efficiently. diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index cec9b57f38..cdb9939926 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -19,8 +19,14 @@ from typeguard import typechecked +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class Lookahead(tf.keras.optimizers.legacy.Optimizer): +class Lookahead(optimizer_class): """This class allows to extend optimizers with the lookahead mechanism. The mechanism is proposed by Michael R. Zhang et.al in the paper diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py index bc29f31e88..60f4cb1d45 100644 --- a/tensorflow_addons/optimizers/novograd.py +++ b/tensorflow_addons/optimizers/novograd.py @@ -21,8 +21,14 @@ from typeguard import typechecked +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class NovoGrad(tf.keras.optimizers.legacy.Optimizer): +class NovoGrad(optimizer_class): """Optimizer that implements NovoGrad. The NovoGrad Optimizer was first proposed in [Stochastic Gradient diff --git a/tensorflow_addons/optimizers/proximal_adagrad.py b/tensorflow_addons/optimizers/proximal_adagrad.py index d62f0669d4..5f24125f97 100644 --- a/tensorflow_addons/optimizers/proximal_adagrad.py +++ b/tensorflow_addons/optimizers/proximal_adagrad.py @@ -22,8 +22,14 @@ from tensorflow_addons.utils.types import FloatTensorLike +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class ProximalAdagrad(tf.keras.optimizers.legacy.Optimizer): +class ProximalAdagrad(optimizer_class): """Optimizer that implements the Proximal Adagrad algorithm. References: diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py index 0c540bfdf1..2fc6253bc6 100644 --- a/tensorflow_addons/optimizers/rectified_adam.py +++ b/tensorflow_addons/optimizers/rectified_adam.py @@ -19,9 +19,14 @@ from typing import Union, Callable, Dict from typeguard import typechecked +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + @tf.keras.utils.register_keras_serializable(package="Addons") -class RectifiedAdam(tf.keras.optimizers.legacy.Optimizer): +class RectifiedAdam(optimizer_class): """Variant of the Adam optimizer whose adaptive learning rate is rectified so as to have a consistent variance. diff --git a/tensorflow_addons/optimizers/yogi.py b/tensorflow_addons/optimizers/yogi.py index 4d602f71f4..a352da61c9 100644 --- a/tensorflow_addons/optimizers/yogi.py +++ b/tensorflow_addons/optimizers/yogi.py @@ -49,8 +49,14 @@ def _solve(a, b, c): return w +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class Yogi(tf.keras.optimizers.legacy.Optimizer): +class Yogi(optimizer_class): """Optimizer that implements the Yogi algorithm in Keras. See Algorithm 2 of From ca68189ecead707ad5914ea8893dcd5ba77824a7 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Wed, 25 May 2022 14:37:35 -0700 Subject: [PATCH 04/14] small fix --- tensorflow_addons/optimizers/conditional_gradient.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tensorflow_addons/optimizers/conditional_gradient.py b/tensorflow_addons/optimizers/conditional_gradient.py index 65d7ae12f4..e3190a6cec 100644 --- a/tensorflow_addons/optimizers/conditional_gradient.py +++ b/tensorflow_addons/optimizers/conditional_gradient.py @@ -21,8 +21,14 @@ from typing import Union, Callable +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @tf.keras.utils.register_keras_serializable(package="Addons") -class ConditionalGradient(tf.keras.optimizers.legacy.Optimizer): +class ConditionalGradient(optimizer_class): """Optimizer that implements the Conditional Gradient optimization. This optimizer helps handle constraints well. From 7c8c291bfaa8fa32bfb49eccbcb1c58ae33073e8 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Thu, 26 May 2022 11:13:17 -0700 Subject: [PATCH 05/14] Update discriminative_layer_training.py --- tensorflow_addons/optimizers/discriminative_layer_training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow_addons/optimizers/discriminative_layer_training.py b/tensorflow_addons/optimizers/discriminative_layer_training.py index 823c3fc0be..8a59b37210 100644 --- a/tensorflow_addons/optimizers/discriminative_layer_training.py +++ b/tensorflow_addons/optimizers/discriminative_layer_training.py @@ -136,7 +136,7 @@ def get_config(self): @classmethod def create_optimizer_spec( cls, - optimizer: tf.keras.optimizers.legacy.Optimizer, + optimizer: optimizer_class, layers_or_model: Union[ tf.keras.Model, tf.keras.Sequential, From 1b16facf65e684153de93eecf97bb30322daac71 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Thu, 26 May 2022 12:15:11 -0700 Subject: [PATCH 06/14] fix version control --- .../optimizers/average_wrapper.py | 22 ++++++++++++------- tensorflow_addons/optimizers/lookahead.py | 22 ++++++++++++------- .../tests/weight_decay_optimizers_test.py | 10 ++++++--- .../optimizers/weight_decay_optimizers.py | 14 ++++++++---- tensorflow_addons/utils/test_utils.py | 1 - 5 files changed, 45 insertions(+), 24 deletions(-) diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index 268138435a..f49ea2f320 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -38,14 +38,20 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if not isinstance( - optimizer, - (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), - ): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " - "or tf.keras.optimizers.Optimizer." - ) + if tf.__version__[:3] > "2.8": + if not isinstance( + optimizer, + (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), + ): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " + "or tf.keras.optimizers.Optimizer." + ) + else: + if not isinstance(optimizer, tf.keras.optimizers.Optimizer): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.Optimizer." + ) self._optimizer = optimizer self._track_trackable(self._optimizer, "awg_optimizer") diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index cdb9939926..a26c5613b6 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -77,14 +77,20 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if not isinstance( - optimizer, - (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), - ): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " - "or tf.keras.optimizers.Optimizer." - ) + if tf.__version__[:3] > "2.8": + if not isinstance( + optimizer, + (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), + ): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " + "or tf.keras.optimizers.Optimizer." + ) + else: + if not isinstance(optimizer, tf.keras.optimizers.Optimizer): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.Optimizer." + ) self._optimizer = optimizer self._set_hyper("sync_period", sync_period) diff --git a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py index 2ba5a3944c..d343045633 100644 --- a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py +++ b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py @@ -401,13 +401,17 @@ def test_var_list_with_exclude_list_sgdw(dtype): ) +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.SGD +else: + optimizer_class = tf.keras.optimizers.SGD + + @pytest.mark.parametrize( "optimizer", [ weight_decay_optimizers.SGDW, - weight_decay_optimizers.extend_with_decoupled_weight_decay( - tf.keras.optimizers.legacy.SGD - ), + weight_decay_optimizers.extend_with_decoupled_weight_decay(optimizer_class), ], ) @pytest.mark.parametrize("dtype", [(tf.half, 0), (tf.float32, 1), (tf.float64, 2)]) diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py index e0405b81ba..39151ea506 100644 --- a/tensorflow_addons/optimizers/weight_decay_optimizers.py +++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py @@ -261,12 +261,18 @@ def _do_use_weight_decay(self, var): return var.ref() in self._decay_var_list +if tf.__version__[:3] > "2.8": + optimizer_class = Union[ + tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer + ] +else: + optimizer_class = tf.keras.optimizers.Optimizer + + @typechecked def extend_with_decoupled_weight_decay( - base_optimizer: Type[ - Union[tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer] - ], -) -> Type[Union[tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer]]: + base_optimizer: Type[optimizer_class], +) -> Type[optimizer_class]: """Factory function returning an optimizer class with decoupled weight decay. diff --git a/tensorflow_addons/utils/test_utils.py b/tensorflow_addons/utils/test_utils.py index e137c19e98..fde88c6a26 100644 --- a/tensorflow_addons/utils/test_utils.py +++ b/tensorflow_addons/utils/test_utils.py @@ -25,7 +25,6 @@ from tensorflow_addons import options from tensorflow_addons.utils import resource_loader - if tf.__version__[:3] > "2.8": from keras.testing_infra.test_utils import layer_test # noqa: F401 else: From 13c58f62aed2b2e9062c23bc57d71f8442922c35 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Thu, 26 May 2022 14:41:34 -0700 Subject: [PATCH 07/14] small fix --- tensorflow_addons/optimizers/tests/standard_test.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorflow_addons/optimizers/tests/standard_test.py b/tensorflow_addons/optimizers/tests/standard_test.py index 948ef6d2b0..107fd106e7 100644 --- a/tensorflow_addons/optimizers/tests/standard_test.py +++ b/tensorflow_addons/optimizers/tests/standard_test.py @@ -31,10 +31,12 @@ "MovingAverage", # is wrapper ] +if tf.__version__[:3] > "2.8": + optimizer_class = tf.keras.optimizers.legacy.Optimizer +else: + optimizer_class = tf.keras.optimizers.Optimizer -classes_to_test = discover_classes( - optimizers, tf.keras.optimizers.legacy.Optimizer, class_exceptions -) +classes_to_test = discover_classes(optimizers, optimizer_class, class_exceptions) @pytest.mark.parametrize("optimizer", classes_to_test) From 743c51df58758808dee6e5dc60aaf7c559ee39b9 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Wed, 1 Jun 2022 12:52:19 -0700 Subject: [PATCH 08/14] move optimizer class to __init__.py --- tensorflow_addons/optimizers/__init__.py | 1 + tensorflow_addons/optimizers/adabelief.py | 8 ++------ tensorflow_addons/optimizers/average_wrapper.py | 10 ++-------- tensorflow_addons/optimizers/cocob.py | 8 ++------ tensorflow_addons/optimizers/conditional_gradient.py | 9 ++------- tensorflow_addons/optimizers/constants.py | 5 +++++ .../optimizers/discriminative_layer_training.py | 12 ++++-------- tensorflow_addons/optimizers/lamb.py | 8 ++------ tensorflow_addons/optimizers/lookahead.py | 9 ++------- tensorflow_addons/optimizers/novograd.py | 10 ++-------- tensorflow_addons/optimizers/proximal_adagrad.py | 9 ++------- tensorflow_addons/optimizers/rectified_adam.py | 8 ++------ tensorflow_addons/optimizers/tests/standard_test.py | 9 +++------ .../optimizers/weight_decay_optimizers.py | 9 ++++----- tensorflow_addons/optimizers/yogi.py | 9 ++------- 15 files changed, 37 insertions(+), 87 deletions(-) create mode 100644 tensorflow_addons/optimizers/constants.py diff --git a/tensorflow_addons/optimizers/__init__.py b/tensorflow_addons/optimizers/__init__.py index 6a23397175..769c3ba866 100644 --- a/tensorflow_addons/optimizers/__init__.py +++ b/tensorflow_addons/optimizers/__init__.py @@ -14,6 +14,7 @@ # ============================================================================== """Additional optimizers that conform to Keras API.""" +from tensorflow_addons.optimizers.constants import BASE_OPTIMIZER_CLASS from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper from tensorflow_addons.optimizers.conditional_gradient import ConditionalGradient from tensorflow_addons.optimizers.cyclical_learning_rate import CyclicalLearningRate diff --git a/tensorflow_addons/optimizers/adabelief.py b/tensorflow_addons/optimizers/adabelief.py index 8c3317c4e9..7dfef441c2 100644 --- a/tensorflow_addons/optimizers/adabelief.py +++ b/tensorflow_addons/optimizers/adabelief.py @@ -17,16 +17,12 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typing import Union, Callable, Dict -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - @tf.keras.utils.register_keras_serializable(package="Addons") -class AdaBelief(optimizer_class): +class AdaBelief(BASE_OPTIMIZER_CLASS): """Variant of the Adam optimizer. It achieves fast convergence as Adam and generalization comparable to SGD. diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index f49ea2f320..ad5a2d93fe 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -17,18 +17,12 @@ import warnings import tensorflow as tf +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from tensorflow_addons.utils import types - from typeguard import typechecked -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - - -class AveragedOptimizerWrapper(optimizer_class, metaclass=abc.ABCMeta): +class AveragedOptimizerWrapper(BASE_OPTIMIZER_CLASS, metaclass=abc.ABCMeta): @typechecked def __init__( self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs diff --git a/tensorflow_addons/optimizers/cocob.py b/tensorflow_addons/optimizers/cocob.py index ad2c3e5178..09f376a979 100644 --- a/tensorflow_addons/optimizers/cocob.py +++ b/tensorflow_addons/optimizers/cocob.py @@ -17,15 +17,11 @@ from typeguard import typechecked import tensorflow as tf - -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS @tf.keras.utils.register_keras_serializable(package="Addons") -class COCOB(optimizer_class): +class COCOB(BASE_OPTIMIZER_CLASS): """Optimizer that implements COCOB Backprop Algorithm Reference: diff --git a/tensorflow_addons/optimizers/conditional_gradient.py b/tensorflow_addons/optimizers/conditional_gradient.py index e3190a6cec..fe2a8253f5 100644 --- a/tensorflow_addons/optimizers/conditional_gradient.py +++ b/tensorflow_addons/optimizers/conditional_gradient.py @@ -17,18 +17,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typeguard import typechecked from typing import Union, Callable -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - - @tf.keras.utils.register_keras_serializable(package="Addons") -class ConditionalGradient(optimizer_class): +class ConditionalGradient(BASE_OPTIMIZER_CLASS): """Optimizer that implements the Conditional Gradient optimization. This optimizer helps handle constraints well. diff --git a/tensorflow_addons/optimizers/constants.py b/tensorflow_addons/optimizers/constants.py new file mode 100644 index 0000000000..3077ce7998 --- /dev/null +++ b/tensorflow_addons/optimizers/constants.py @@ -0,0 +1,5 @@ +import tensorflow as tf + +BASE_OPTIMIZER_CLASS = tf.keras.optimizers.legacy.Optimizer +if tf.__version__[:3] <= "2.8": + BASE_OPTIMIZER_CLASS = tf.keras.optimizers.Optimizer diff --git a/tensorflow_addons/optimizers/discriminative_layer_training.py b/tensorflow_addons/optimizers/discriminative_layer_training.py index 8a59b37210..54decbd875 100644 --- a/tensorflow_addons/optimizers/discriminative_layer_training.py +++ b/tensorflow_addons/optimizers/discriminative_layer_training.py @@ -17,17 +17,13 @@ from typing import List, Union import tensorflow as tf -from typeguard import typechecked - -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class MultiOptimizer(optimizer_class): +class MultiOptimizer(BASE_OPTIMIZER_CLASS): """Multi Optimizer Wrapper for Discriminative Layer Training. Creates a wrapper around a set of instantiated optimizer layer pairs. @@ -136,7 +132,7 @@ def get_config(self): @classmethod def create_optimizer_spec( cls, - optimizer: optimizer_class, + optimizer: BASE_OPTIMIZER_CLASS, layers_or_model: Union[ tf.keras.Model, tf.keras.Sequential, diff --git a/tensorflow_addons/optimizers/lamb.py b/tensorflow_addons/optimizers/lamb.py index d96129498a..9ff6ce8fe8 100644 --- a/tensorflow_addons/optimizers/lamb.py +++ b/tensorflow_addons/optimizers/lamb.py @@ -24,17 +24,13 @@ from typeguard import typechecked import tensorflow as tf +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from tensorflow_addons.utils.types import FloatTensorLike from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - @tf.keras.utils.register_keras_serializable(package="Addons") -class LAMB(optimizer_class): +class LAMB(BASE_OPTIMIZER_CLASS): """Optimizer that implements the Layer-wise Adaptive Moments (LAMB). See paper [Large Batch Optimization for Deep Learning: Training BERT diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index a26c5613b6..495a5d788f 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -16,17 +16,12 @@ import tensorflow as tf from tensorflow_addons.utils import types +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typeguard import typechecked -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - - @tf.keras.utils.register_keras_serializable(package="Addons") -class Lookahead(optimizer_class): +class Lookahead(BASE_OPTIMIZER_CLASS): """This class allows to extend optimizers with the lookahead mechanism. The mechanism is proposed by Michael R. Zhang et.al in the paper diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py index 60f4cb1d45..6c9aa2ebbe 100644 --- a/tensorflow_addons/optimizers/novograd.py +++ b/tensorflow_addons/optimizers/novograd.py @@ -16,19 +16,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike - +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typing import Union, Callable from typeguard import typechecked -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - - @tf.keras.utils.register_keras_serializable(package="Addons") -class NovoGrad(optimizer_class): +class NovoGrad(BASE_OPTIMIZER_CLASS): """Optimizer that implements NovoGrad. The NovoGrad Optimizer was first proposed in [Stochastic Gradient diff --git a/tensorflow_addons/optimizers/proximal_adagrad.py b/tensorflow_addons/optimizers/proximal_adagrad.py index 5f24125f97..01df4315aa 100644 --- a/tensorflow_addons/optimizers/proximal_adagrad.py +++ b/tensorflow_addons/optimizers/proximal_adagrad.py @@ -19,17 +19,12 @@ import tensorflow as tf from typeguard import typechecked +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from tensorflow_addons.utils.types import FloatTensorLike -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - - @tf.keras.utils.register_keras_serializable(package="Addons") -class ProximalAdagrad(optimizer_class): +class ProximalAdagrad(BASE_OPTIMIZER_CLASS): """Optimizer that implements the Proximal Adagrad algorithm. References: diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py index 2fc6253bc6..3712d9fb5a 100644 --- a/tensorflow_addons/optimizers/rectified_adam.py +++ b/tensorflow_addons/optimizers/rectified_adam.py @@ -16,17 +16,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typing import Union, Callable, Dict from typeguard import typechecked -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - @tf.keras.utils.register_keras_serializable(package="Addons") -class RectifiedAdam(optimizer_class): +class RectifiedAdam(BASE_OPTIMIZER_CLASS): """Variant of the Adam optimizer whose adaptive learning rate is rectified so as to have a consistent variance. diff --git a/tensorflow_addons/optimizers/tests/standard_test.py b/tensorflow_addons/optimizers/tests/standard_test.py index 107fd106e7..f48fc4f28d 100644 --- a/tensorflow_addons/optimizers/tests/standard_test.py +++ b/tensorflow_addons/optimizers/tests/standard_test.py @@ -18,6 +18,7 @@ import tensorflow as tf from tensorflow_addons import optimizers +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from tensorflow_addons.utils.test_utils import discover_classes class_exceptions = [ @@ -29,14 +30,10 @@ "ConditionalGradient", # is wrapper "Lookahead", # is wrapper "MovingAverage", # is wrapper + "BASE_OPTIMIZER_CLASS", # is a constantc ] -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - -classes_to_test = discover_classes(optimizers, optimizer_class, class_exceptions) +classes_to_test = discover_classes(optimizers, BASE_OPTIMIZER_CLASS, class_exceptions) @pytest.mark.parametrize("optimizer", classes_to_test) diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py index 39151ea506..91279524ab 100644 --- a/tensorflow_addons/optimizers/weight_decay_optimizers.py +++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py @@ -261,11 +261,10 @@ def _do_use_weight_decay(self, var): return var.ref() in self._decay_var_list -if tf.__version__[:3] > "2.8": - optimizer_class = Union[ - tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer - ] -else: +optimizer_class = Union[ + tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer +] +if tf.__version__[:3] <= "2.8": optimizer_class = tf.keras.optimizers.Optimizer diff --git a/tensorflow_addons/optimizers/yogi.py b/tensorflow_addons/optimizers/yogi.py index a352da61c9..057520192d 100644 --- a/tensorflow_addons/optimizers/yogi.py +++ b/tensorflow_addons/optimizers/yogi.py @@ -25,6 +25,7 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typeguard import typechecked from typing import Union, Callable @@ -49,14 +50,8 @@ def _solve(a, b, c): return w -if tf.__version__[:3] > "2.8": - optimizer_class = tf.keras.optimizers.legacy.Optimizer -else: - optimizer_class = tf.keras.optimizers.Optimizer - - @tf.keras.utils.register_keras_serializable(package="Addons") -class Yogi(optimizer_class): +class Yogi(BASE_OPTIMIZER_CLASS): """Optimizer that implements the Yogi algorithm in Keras. See Algorithm 2 of From a66a5956e85122a754c0f6e3de8de0dcdf5cb3aa Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Wed, 1 Jun 2022 12:56:41 -0700 Subject: [PATCH 09/14] small fix --- tensorflow_addons/optimizers/average_wrapper.py | 12 ++++++------ tensorflow_addons/optimizers/lookahead.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index ad5a2d93fe..8647fd5548 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -32,7 +32,12 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if tf.__version__[:3] > "2.8": + if tf.__version__[:3] <= "2.8": + if not isinstance(optimizer, tf.keras.optimizers.Optimizer): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.Optimizer." + ) + else: if not isinstance( optimizer, (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), @@ -41,11 +46,6 @@ def __init__( "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " "or tf.keras.optimizers.Optimizer." ) - else: - if not isinstance(optimizer, tf.keras.optimizers.Optimizer): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.Optimizer." - ) self._optimizer = optimizer self._track_trackable(self._optimizer, "awg_optimizer") diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index 495a5d788f..149738b1db 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -72,7 +72,12 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if tf.__version__[:3] > "2.8": + if tf.__version__[:3] <= "2.8": + if not isinstance(optimizer, tf.keras.optimizers.Optimizer): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.Optimizer." + ) + else: if not isinstance( optimizer, (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), @@ -81,11 +86,6 @@ def __init__( "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " "or tf.keras.optimizers.Optimizer." ) - else: - if not isinstance(optimizer, tf.keras.optimizers.Optimizer): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.Optimizer." - ) self._optimizer = optimizer self._set_hyper("sync_period", sync_period) From 69bcdb43a411d3ab56875b2ffbfce5bfb1f7442b Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Thu, 2 Jun 2022 13:42:06 -0700 Subject: [PATCH 10/14] fix problems --- .../optimizers/average_wrapper.py | 21 +++++++------------ .../optimizers/conditional_gradient.py | 2 +- tensorflow_addons/optimizers/constants.py | 20 ++++++++++++++++-- tensorflow_addons/optimizers/lookahead.py | 21 +++++++------------ tensorflow_addons/utils/types.py | 8 ++++++- 5 files changed, 40 insertions(+), 32 deletions(-) diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index 8647fd5548..f7c17668a3 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -32,20 +32,13 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if tf.__version__[:3] <= "2.8": - if not isinstance(optimizer, tf.keras.optimizers.Optimizer): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.Optimizer." - ) - else: - if not isinstance( - optimizer, - (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), - ): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " - "or tf.keras.optimizers.Optimizer." - ) + if not isinstance( + optimizer, (tf.keras.optimizers.Optimizer, BASE_OPTIMIZER_CLASS) + ): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.Optimizer " + "or tf.keras.optimizers.legacy.Optimizer (if you have tf version >= 2.9.0)." + ) self._optimizer = optimizer self._track_trackable(self._optimizer, "awg_optimizer") diff --git a/tensorflow_addons/optimizers/conditional_gradient.py b/tensorflow_addons/optimizers/conditional_gradient.py index fe2a8253f5..37d8c4a480 100644 --- a/tensorflow_addons/optimizers/conditional_gradient.py +++ b/tensorflow_addons/optimizers/conditional_gradient.py @@ -15,9 +15,9 @@ """Conditional Gradient optimizer.""" import tensorflow as tf +from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS from typeguard import typechecked from typing import Union, Callable diff --git a/tensorflow_addons/optimizers/constants.py b/tensorflow_addons/optimizers/constants.py index 3077ce7998..00eb160685 100644 --- a/tensorflow_addons/optimizers/constants.py +++ b/tensorflow_addons/optimizers/constants.py @@ -1,5 +1,21 @@ +# Copyright 2022 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import importlib import tensorflow as tf -BASE_OPTIMIZER_CLASS = tf.keras.optimizers.legacy.Optimizer -if tf.__version__[:3] <= "2.8": +if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: + BASE_OPTIMIZER_CLASS = tf.keras.optimizers.legacy.Optimizer +else: BASE_OPTIMIZER_CLASS = tf.keras.optimizers.Optimizer diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index 149738b1db..7c3266d997 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -72,20 +72,13 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) - if tf.__version__[:3] <= "2.8": - if not isinstance(optimizer, tf.keras.optimizers.Optimizer): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.Optimizer." - ) - else: - if not isinstance( - optimizer, - (tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer), - ): - raise TypeError( - "optimizer is not an object of tf.keras.optimizers.legacy.Optimizer " - "or tf.keras.optimizers.Optimizer." - ) + if not isinstance( + optimizer, (tf.keras.optimizers.Optimizer, BASE_OPTIMIZER_CLASS) + ): + raise TypeError( + "optimizer is not an object of tf.keras.optimizers.Optimizer " + "or tf.keras.optimizers.legacy.Optimizer (if you have tf version >= 2.9.0)." + ) self._optimizer = optimizer self._set_hyper("sync_period", sync_period) diff --git a/tensorflow_addons/utils/types.py b/tensorflow_addons/utils/types.py index 1c2cb0851d..4bfa0dacf6 100644 --- a/tensorflow_addons/utils/types.py +++ b/tensorflow_addons/utils/types.py @@ -16,6 +16,7 @@ from typing import Union, Callable, List +import importlib import numpy as np import tensorflow as tf @@ -46,7 +47,12 @@ Regularizer = Union[None, dict, str, Callable, tf.keras.regularizers.Regularizer] Constraint = Union[None, dict, str, Callable, tf.keras.constraints.Constraint] Activation = Union[None, str, Callable] -Optimizer = Union[tf.keras.optimizers.Optimizer, str] +if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: + Optimizer = Union[ + tf.keras.optimizers.Optimizer, tf.keras.optimizers.legacy.Optimizer, str + ] +else: + Optimizer = Union[tf.keras.optimizers.Optimizer, str] TensorLike = Union[ List[Union[Number, list]], From e734793bcc051b8605a1c97f8da37796d49e3bd0 Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Thu, 2 Jun 2022 13:58:42 -0700 Subject: [PATCH 11/14] small fix --- tensorflow_addons/optimizers/__init__.py | 2 +- tensorflow_addons/optimizers/adabelief.py | 4 ++-- tensorflow_addons/optimizers/average_wrapper.py | 6 +++--- tensorflow_addons/optimizers/cocob.py | 4 ++-- tensorflow_addons/optimizers/conditional_gradient.py | 4 ++-- tensorflow_addons/optimizers/constants.py | 4 ++-- .../optimizers/discriminative_layer_training.py | 6 +++--- tensorflow_addons/optimizers/lamb.py | 4 ++-- tensorflow_addons/optimizers/lookahead.py | 6 +++--- tensorflow_addons/optimizers/novograd.py | 4 ++-- tensorflow_addons/optimizers/proximal_adagrad.py | 4 ++-- tensorflow_addons/optimizers/rectified_adam.py | 4 ++-- tensorflow_addons/optimizers/tests/standard_test.py | 6 +++--- tensorflow_addons/optimizers/yogi.py | 4 ++-- 14 files changed, 31 insertions(+), 31 deletions(-) diff --git a/tensorflow_addons/optimizers/__init__.py b/tensorflow_addons/optimizers/__init__.py index 769c3ba866..6b80b07864 100644 --- a/tensorflow_addons/optimizers/__init__.py +++ b/tensorflow_addons/optimizers/__init__.py @@ -14,7 +14,7 @@ # ============================================================================== """Additional optimizers that conform to Keras API.""" -from tensorflow_addons.optimizers.constants import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers.constants import BaseOptimizerClass from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper from tensorflow_addons.optimizers.conditional_gradient import ConditionalGradient from tensorflow_addons.optimizers.cyclical_learning_rate import CyclicalLearningRate diff --git a/tensorflow_addons/optimizers/adabelief.py b/tensorflow_addons/optimizers/adabelief.py index 7dfef441c2..c58a1c40d5 100644 --- a/tensorflow_addons/optimizers/adabelief.py +++ b/tensorflow_addons/optimizers/adabelief.py @@ -17,12 +17,12 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from typing import Union, Callable, Dict @tf.keras.utils.register_keras_serializable(package="Addons") -class AdaBelief(BASE_OPTIMIZER_CLASS): +class AdaBelief(BaseOptimizerClass): """Variant of the Adam optimizer. It achieves fast convergence as Adam and generalization comparable to SGD. diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index f7c17668a3..c6749ad8e4 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -17,12 +17,12 @@ import warnings import tensorflow as tf -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from tensorflow_addons.utils import types from typeguard import typechecked -class AveragedOptimizerWrapper(BASE_OPTIMIZER_CLASS, metaclass=abc.ABCMeta): +class AveragedOptimizerWrapper(BaseOptimizerClass, metaclass=abc.ABCMeta): @typechecked def __init__( self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs @@ -33,7 +33,7 @@ def __init__( optimizer = tf.keras.optimizers.get(optimizer) if not isinstance( - optimizer, (tf.keras.optimizers.Optimizer, BASE_OPTIMIZER_CLASS) + optimizer, (tf.keras.optimizers.Optimizer, BaseOptimizerClass) ): raise TypeError( "optimizer is not an object of tf.keras.optimizers.Optimizer " diff --git a/tensorflow_addons/optimizers/cocob.py b/tensorflow_addons/optimizers/cocob.py index 09f376a979..194d37f10d 100644 --- a/tensorflow_addons/optimizers/cocob.py +++ b/tensorflow_addons/optimizers/cocob.py @@ -17,11 +17,11 @@ from typeguard import typechecked import tensorflow as tf -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass @tf.keras.utils.register_keras_serializable(package="Addons") -class COCOB(BASE_OPTIMIZER_CLASS): +class COCOB(BaseOptimizerClass): """Optimizer that implements COCOB Backprop Algorithm Reference: diff --git a/tensorflow_addons/optimizers/conditional_gradient.py b/tensorflow_addons/optimizers/conditional_gradient.py index 37d8c4a480..ccedd7465a 100644 --- a/tensorflow_addons/optimizers/conditional_gradient.py +++ b/tensorflow_addons/optimizers/conditional_gradient.py @@ -15,7 +15,7 @@ """Conditional Gradient optimizer.""" import tensorflow as tf -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from tensorflow_addons.utils.types import FloatTensorLike from typeguard import typechecked @@ -23,7 +23,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class ConditionalGradient(BASE_OPTIMIZER_CLASS): +class ConditionalGradient(BaseOptimizerClass): """Optimizer that implements the Conditional Gradient optimization. This optimizer helps handle constraints well. diff --git a/tensorflow_addons/optimizers/constants.py b/tensorflow_addons/optimizers/constants.py index 00eb160685..e8da5f3422 100644 --- a/tensorflow_addons/optimizers/constants.py +++ b/tensorflow_addons/optimizers/constants.py @@ -16,6 +16,6 @@ import tensorflow as tf if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: - BASE_OPTIMIZER_CLASS = tf.keras.optimizers.legacy.Optimizer + BaseOptimizerClass = tf.keras.optimizers.legacy.Optimizer else: - BASE_OPTIMIZER_CLASS = tf.keras.optimizers.Optimizer + BaseOptimizerClass = tf.keras.optimizers.Optimizer diff --git a/tensorflow_addons/optimizers/discriminative_layer_training.py b/tensorflow_addons/optimizers/discriminative_layer_training.py index 54decbd875..e7acacef72 100644 --- a/tensorflow_addons/optimizers/discriminative_layer_training.py +++ b/tensorflow_addons/optimizers/discriminative_layer_training.py @@ -18,12 +18,12 @@ import tensorflow as tf -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class MultiOptimizer(BASE_OPTIMIZER_CLASS): +class MultiOptimizer(BaseOptimizerClass): """Multi Optimizer Wrapper for Discriminative Layer Training. Creates a wrapper around a set of instantiated optimizer layer pairs. @@ -132,7 +132,7 @@ def get_config(self): @classmethod def create_optimizer_spec( cls, - optimizer: BASE_OPTIMIZER_CLASS, + optimizer: BaseOptimizerClass, layers_or_model: Union[ tf.keras.Model, tf.keras.Sequential, diff --git a/tensorflow_addons/optimizers/lamb.py b/tensorflow_addons/optimizers/lamb.py index 9ff6ce8fe8..c56a9693ca 100644 --- a/tensorflow_addons/optimizers/lamb.py +++ b/tensorflow_addons/optimizers/lamb.py @@ -24,13 +24,13 @@ from typeguard import typechecked import tensorflow as tf -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from tensorflow_addons.utils.types import FloatTensorLike from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes @tf.keras.utils.register_keras_serializable(package="Addons") -class LAMB(BASE_OPTIMIZER_CLASS): +class LAMB(BaseOptimizerClass): """Optimizer that implements the Layer-wise Adaptive Moments (LAMB). See paper [Large Batch Optimization for Deep Learning: Training BERT diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index 7c3266d997..5a27953b21 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -16,12 +16,12 @@ import tensorflow as tf from tensorflow_addons.utils import types -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class Lookahead(BASE_OPTIMIZER_CLASS): +class Lookahead(BaseOptimizerClass): """This class allows to extend optimizers with the lookahead mechanism. The mechanism is proposed by Michael R. Zhang et.al in the paper @@ -73,7 +73,7 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) if not isinstance( - optimizer, (tf.keras.optimizers.Optimizer, BASE_OPTIMIZER_CLASS) + optimizer, (tf.keras.optimizers.Optimizer, BaseOptimizerClass) ): raise TypeError( "optimizer is not an object of tf.keras.optimizers.Optimizer " diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py index 6c9aa2ebbe..ea7a8a34da 100644 --- a/tensorflow_addons/optimizers/novograd.py +++ b/tensorflow_addons/optimizers/novograd.py @@ -16,13 +16,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from typing import Union, Callable from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class NovoGrad(BASE_OPTIMIZER_CLASS): +class NovoGrad(BaseOptimizerClass): """Optimizer that implements NovoGrad. The NovoGrad Optimizer was first proposed in [Stochastic Gradient diff --git a/tensorflow_addons/optimizers/proximal_adagrad.py b/tensorflow_addons/optimizers/proximal_adagrad.py index 01df4315aa..f370b08599 100644 --- a/tensorflow_addons/optimizers/proximal_adagrad.py +++ b/tensorflow_addons/optimizers/proximal_adagrad.py @@ -19,12 +19,12 @@ import tensorflow as tf from typeguard import typechecked -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from tensorflow_addons.utils.types import FloatTensorLike @tf.keras.utils.register_keras_serializable(package="Addons") -class ProximalAdagrad(BASE_OPTIMIZER_CLASS): +class ProximalAdagrad(BaseOptimizerClass): """Optimizer that implements the Proximal Adagrad algorithm. References: diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py index 3712d9fb5a..053c73d185 100644 --- a/tensorflow_addons/optimizers/rectified_adam.py +++ b/tensorflow_addons/optimizers/rectified_adam.py @@ -16,13 +16,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from typing import Union, Callable, Dict from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class RectifiedAdam(BASE_OPTIMIZER_CLASS): +class RectifiedAdam(BaseOptimizerClass): """Variant of the Adam optimizer whose adaptive learning rate is rectified so as to have a consistent variance. diff --git a/tensorflow_addons/optimizers/tests/standard_test.py b/tensorflow_addons/optimizers/tests/standard_test.py index f48fc4f28d..763dfba08d 100644 --- a/tensorflow_addons/optimizers/tests/standard_test.py +++ b/tensorflow_addons/optimizers/tests/standard_test.py @@ -18,7 +18,7 @@ import tensorflow as tf from tensorflow_addons import optimizers -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from tensorflow_addons.utils.test_utils import discover_classes class_exceptions = [ @@ -30,10 +30,10 @@ "ConditionalGradient", # is wrapper "Lookahead", # is wrapper "MovingAverage", # is wrapper - "BASE_OPTIMIZER_CLASS", # is a constantc + "BaseOptimizerClass", # is a constantc ] -classes_to_test = discover_classes(optimizers, BASE_OPTIMIZER_CLASS, class_exceptions) +classes_to_test = discover_classes(optimizers, BaseOptimizerClass, class_exceptions) @pytest.mark.parametrize("optimizer", classes_to_test) diff --git a/tensorflow_addons/optimizers/yogi.py b/tensorflow_addons/optimizers/yogi.py index 057520192d..25fc07852f 100644 --- a/tensorflow_addons/optimizers/yogi.py +++ b/tensorflow_addons/optimizers/yogi.py @@ -25,7 +25,7 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BASE_OPTIMIZER_CLASS +from tensorflow_addons.optimizers import BaseOptimizerClass from typeguard import typechecked from typing import Union, Callable @@ -51,7 +51,7 @@ def _solve(a, b, c): @tf.keras.utils.register_keras_serializable(package="Addons") -class Yogi(BASE_OPTIMIZER_CLASS): +class Yogi(BaseOptimizerClass): """Optimizer that implements the Yogi algorithm in Keras. See Algorithm 2 of From 8ecfab532734fa7dc2ae583f23486bf5eaa7f2d6 Mon Sep 17 00:00:00 2001 From: chenmoneygithub Date: Fri, 10 Jun 2022 11:49:59 -0700 Subject: [PATCH 12/14] Rename BaseOptimizer to KerasLegacyOptimizer --- tensorflow_addons/optimizers/__init__.py | 2 +- tensorflow_addons/optimizers/adabelief.py | 4 ++-- tensorflow_addons/optimizers/average_wrapper.py | 6 +++--- tensorflow_addons/optimizers/cocob.py | 4 ++-- .../optimizers/conditional_gradient.py | 4 ++-- tensorflow_addons/optimizers/constants.py | 4 ++-- .../optimizers/discriminative_layer_training.py | 6 +++--- tensorflow_addons/optimizers/lamb.py | 4 ++-- tensorflow_addons/optimizers/lazy_adam.py | 3 ++- tensorflow_addons/optimizers/lookahead.py | 6 +++--- tensorflow_addons/optimizers/novograd.py | 4 ++-- tensorflow_addons/optimizers/proximal_adagrad.py | 4 ++-- tensorflow_addons/optimizers/rectified_adam.py | 4 ++-- .../optimizers/tests/standard_test.py | 6 +++--- .../tests/weight_decay_optimizers_test.py | 3 ++- .../optimizers/weight_decay_optimizers.py | 16 +++++++++------- tensorflow_addons/optimizers/yogi.py | 4 ++-- 17 files changed, 44 insertions(+), 40 deletions(-) diff --git a/tensorflow_addons/optimizers/__init__.py b/tensorflow_addons/optimizers/__init__.py index 6b80b07864..07b52a83fa 100644 --- a/tensorflow_addons/optimizers/__init__.py +++ b/tensorflow_addons/optimizers/__init__.py @@ -14,7 +14,7 @@ # ============================================================================== """Additional optimizers that conform to Keras API.""" -from tensorflow_addons.optimizers.constants import BaseOptimizerClass +from tensorflow_addons.optimizers.constants import KerasLegacyOptimizer from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper from tensorflow_addons.optimizers.conditional_gradient import ConditionalGradient from tensorflow_addons.optimizers.cyclical_learning_rate import CyclicalLearningRate diff --git a/tensorflow_addons/optimizers/adabelief.py b/tensorflow_addons/optimizers/adabelief.py index c58a1c40d5..58808ed696 100644 --- a/tensorflow_addons/optimizers/adabelief.py +++ b/tensorflow_addons/optimizers/adabelief.py @@ -17,12 +17,12 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from typing import Union, Callable, Dict @tf.keras.utils.register_keras_serializable(package="Addons") -class AdaBelief(BaseOptimizerClass): +class AdaBelief(KerasLegacyOptimizer): """Variant of the Adam optimizer. It achieves fast convergence as Adam and generalization comparable to SGD. diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py index c6749ad8e4..fe67b9bd74 100644 --- a/tensorflow_addons/optimizers/average_wrapper.py +++ b/tensorflow_addons/optimizers/average_wrapper.py @@ -17,12 +17,12 @@ import warnings import tensorflow as tf -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from tensorflow_addons.utils import types from typeguard import typechecked -class AveragedOptimizerWrapper(BaseOptimizerClass, metaclass=abc.ABCMeta): +class AveragedOptimizerWrapper(KerasLegacyOptimizer, metaclass=abc.ABCMeta): @typechecked def __init__( self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs @@ -33,7 +33,7 @@ def __init__( optimizer = tf.keras.optimizers.get(optimizer) if not isinstance( - optimizer, (tf.keras.optimizers.Optimizer, BaseOptimizerClass) + optimizer, (tf.keras.optimizers.Optimizer, KerasLegacyOptimizer) ): raise TypeError( "optimizer is not an object of tf.keras.optimizers.Optimizer " diff --git a/tensorflow_addons/optimizers/cocob.py b/tensorflow_addons/optimizers/cocob.py index 194d37f10d..6e887f06e5 100644 --- a/tensorflow_addons/optimizers/cocob.py +++ b/tensorflow_addons/optimizers/cocob.py @@ -17,11 +17,11 @@ from typeguard import typechecked import tensorflow as tf -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer @tf.keras.utils.register_keras_serializable(package="Addons") -class COCOB(BaseOptimizerClass): +class COCOB(KerasLegacyOptimizer): """Optimizer that implements COCOB Backprop Algorithm Reference: diff --git a/tensorflow_addons/optimizers/conditional_gradient.py b/tensorflow_addons/optimizers/conditional_gradient.py index ccedd7465a..6b79626710 100644 --- a/tensorflow_addons/optimizers/conditional_gradient.py +++ b/tensorflow_addons/optimizers/conditional_gradient.py @@ -15,7 +15,7 @@ """Conditional Gradient optimizer.""" import tensorflow as tf -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from tensorflow_addons.utils.types import FloatTensorLike from typeguard import typechecked @@ -23,7 +23,7 @@ @tf.keras.utils.register_keras_serializable(package="Addons") -class ConditionalGradient(BaseOptimizerClass): +class ConditionalGradient(KerasLegacyOptimizer): """Optimizer that implements the Conditional Gradient optimization. This optimizer helps handle constraints well. diff --git a/tensorflow_addons/optimizers/constants.py b/tensorflow_addons/optimizers/constants.py index e8da5f3422..de6a498b69 100644 --- a/tensorflow_addons/optimizers/constants.py +++ b/tensorflow_addons/optimizers/constants.py @@ -16,6 +16,6 @@ import tensorflow as tf if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: - BaseOptimizerClass = tf.keras.optimizers.legacy.Optimizer + KerasLegacyOptimizer = tf.keras.optimizers.legacy.Optimizer else: - BaseOptimizerClass = tf.keras.optimizers.Optimizer + KerasLegacyOptimizer = tf.keras.optimizers.Optimizer diff --git a/tensorflow_addons/optimizers/discriminative_layer_training.py b/tensorflow_addons/optimizers/discriminative_layer_training.py index e7acacef72..1e51fb6616 100644 --- a/tensorflow_addons/optimizers/discriminative_layer_training.py +++ b/tensorflow_addons/optimizers/discriminative_layer_training.py @@ -18,12 +18,12 @@ import tensorflow as tf -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class MultiOptimizer(BaseOptimizerClass): +class MultiOptimizer(KerasLegacyOptimizer): """Multi Optimizer Wrapper for Discriminative Layer Training. Creates a wrapper around a set of instantiated optimizer layer pairs. @@ -132,7 +132,7 @@ def get_config(self): @classmethod def create_optimizer_spec( cls, - optimizer: BaseOptimizerClass, + optimizer: KerasLegacyOptimizer, layers_or_model: Union[ tf.keras.Model, tf.keras.Sequential, diff --git a/tensorflow_addons/optimizers/lamb.py b/tensorflow_addons/optimizers/lamb.py index c56a9693ca..df06c01d8a 100644 --- a/tensorflow_addons/optimizers/lamb.py +++ b/tensorflow_addons/optimizers/lamb.py @@ -24,13 +24,13 @@ from typeguard import typechecked import tensorflow as tf -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from tensorflow_addons.utils.types import FloatTensorLike from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes @tf.keras.utils.register_keras_serializable(package="Addons") -class LAMB(BaseOptimizerClass): +class LAMB(KerasLegacyOptimizer): """Optimizer that implements the Layer-wise Adaptive Moments (LAMB). See paper [Large Batch Optimization for Deep Learning: Training BERT diff --git a/tensorflow_addons/optimizers/lazy_adam.py b/tensorflow_addons/optimizers/lazy_adam.py index aac7a38706..ad8570bc3c 100644 --- a/tensorflow_addons/optimizers/lazy_adam.py +++ b/tensorflow_addons/optimizers/lazy_adam.py @@ -20,6 +20,7 @@ original Adam algorithm, and may lead to different empirical results. """ +import importlib import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike @@ -27,7 +28,7 @@ from typing import Union, Callable -if tf.__version__[:3] > "2.8": +if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: adam_optimizer_class = tf.keras.optimizers.legacy.Adam else: adam_optimizer_class = tf.keras.optimizers.Adam diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py index 5a27953b21..a10091164d 100644 --- a/tensorflow_addons/optimizers/lookahead.py +++ b/tensorflow_addons/optimizers/lookahead.py @@ -16,12 +16,12 @@ import tensorflow as tf from tensorflow_addons.utils import types -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class Lookahead(BaseOptimizerClass): +class Lookahead(KerasLegacyOptimizer): """This class allows to extend optimizers with the lookahead mechanism. The mechanism is proposed by Michael R. Zhang et.al in the paper @@ -73,7 +73,7 @@ def __init__( if isinstance(optimizer, str): optimizer = tf.keras.optimizers.get(optimizer) if not isinstance( - optimizer, (tf.keras.optimizers.Optimizer, BaseOptimizerClass) + optimizer, (tf.keras.optimizers.Optimizer, KerasLegacyOptimizer) ): raise TypeError( "optimizer is not an object of tf.keras.optimizers.Optimizer " diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py index ea7a8a34da..cea891faac 100644 --- a/tensorflow_addons/optimizers/novograd.py +++ b/tensorflow_addons/optimizers/novograd.py @@ -16,13 +16,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from typing import Union, Callable from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class NovoGrad(BaseOptimizerClass): +class NovoGrad(KerasLegacyOptimizer): """Optimizer that implements NovoGrad. The NovoGrad Optimizer was first proposed in [Stochastic Gradient diff --git a/tensorflow_addons/optimizers/proximal_adagrad.py b/tensorflow_addons/optimizers/proximal_adagrad.py index f370b08599..eb5853cdcb 100644 --- a/tensorflow_addons/optimizers/proximal_adagrad.py +++ b/tensorflow_addons/optimizers/proximal_adagrad.py @@ -19,12 +19,12 @@ import tensorflow as tf from typeguard import typechecked -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from tensorflow_addons.utils.types import FloatTensorLike @tf.keras.utils.register_keras_serializable(package="Addons") -class ProximalAdagrad(BaseOptimizerClass): +class ProximalAdagrad(KerasLegacyOptimizer): """Optimizer that implements the Proximal Adagrad algorithm. References: diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py index 053c73d185..9d963e8014 100644 --- a/tensorflow_addons/optimizers/rectified_adam.py +++ b/tensorflow_addons/optimizers/rectified_adam.py @@ -16,13 +16,13 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from typing import Union, Callable, Dict from typeguard import typechecked @tf.keras.utils.register_keras_serializable(package="Addons") -class RectifiedAdam(BaseOptimizerClass): +class RectifiedAdam(KerasLegacyOptimizer): """Variant of the Adam optimizer whose adaptive learning rate is rectified so as to have a consistent variance. diff --git a/tensorflow_addons/optimizers/tests/standard_test.py b/tensorflow_addons/optimizers/tests/standard_test.py index 763dfba08d..c0568b82f9 100644 --- a/tensorflow_addons/optimizers/tests/standard_test.py +++ b/tensorflow_addons/optimizers/tests/standard_test.py @@ -18,7 +18,7 @@ import tensorflow as tf from tensorflow_addons import optimizers -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from tensorflow_addons.utils.test_utils import discover_classes class_exceptions = [ @@ -30,10 +30,10 @@ "ConditionalGradient", # is wrapper "Lookahead", # is wrapper "MovingAverage", # is wrapper - "BaseOptimizerClass", # is a constantc + "KerasLegacyOptimizer", # is a constantc ] -classes_to_test = discover_classes(optimizers, BaseOptimizerClass, class_exceptions) +classes_to_test = discover_classes(optimizers, KerasLegacyOptimizer, class_exceptions) @pytest.mark.parametrize("optimizer", classes_to_test) diff --git a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py index d343045633..8099c556c0 100644 --- a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py +++ b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py @@ -14,6 +14,7 @@ # ============================================================================== """Tests for optimizers with weight decay.""" +import importlib import numpy as np import pytest import tensorflow as tf @@ -401,7 +402,7 @@ def test_var_list_with_exclude_list_sgdw(dtype): ) -if tf.__version__[:3] > "2.8": +if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: optimizer_class = tf.keras.optimizers.legacy.SGD else: optimizer_class = tf.keras.optimizers.SGD diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py index 91279524ab..1fc491d547 100644 --- a/tensorflow_addons/optimizers/weight_decay_optimizers.py +++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py @@ -14,6 +14,7 @@ # ============================================================================== """Base class to make optimizers weight decay ready.""" +import importlib import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes @@ -261,17 +262,18 @@ def _do_use_weight_decay(self, var): return var.ref() in self._decay_var_list -optimizer_class = Union[ - tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer -] -if tf.__version__[:3] <= "2.8": - optimizer_class = tf.keras.optimizers.Optimizer +if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: + keras_legacy_optimizer = Union[ + tf.keras.optimizers.legacy.Optimizer, tf.keras.optimizers.Optimizer + ] +else: + keras_legacy_optimizer = tf.keras.optimizers.Optimizer @typechecked def extend_with_decoupled_weight_decay( - base_optimizer: Type[optimizer_class], -) -> Type[optimizer_class]: + base_optimizer: Type[keras_legacy_optimizer], +) -> Type[keras_legacy_optimizer]: """Factory function returning an optimizer class with decoupled weight decay. diff --git a/tensorflow_addons/optimizers/yogi.py b/tensorflow_addons/optimizers/yogi.py index 25fc07852f..8d7e0a272b 100644 --- a/tensorflow_addons/optimizers/yogi.py +++ b/tensorflow_addons/optimizers/yogi.py @@ -25,7 +25,7 @@ import tensorflow as tf from tensorflow_addons.utils.types import FloatTensorLike -from tensorflow_addons.optimizers import BaseOptimizerClass +from tensorflow_addons.optimizers import KerasLegacyOptimizer from typeguard import typechecked from typing import Union, Callable @@ -51,7 +51,7 @@ def _solve(a, b, c): @tf.keras.utils.register_keras_serializable(package="Addons") -class Yogi(BaseOptimizerClass): +class Yogi(KerasLegacyOptimizer): """Optimizer that implements the Yogi algorithm in Keras. See Algorithm 2 of From 82554de0994d094ec8541d6e898398957cc775b1 Mon Sep 17 00:00:00 2001 From: chenmoneygithub Date: Sat, 18 Jun 2022 16:41:23 -0700 Subject: [PATCH 13/14] exclude keras optimizer from type check --- tools/testing/source_code_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/testing/source_code_test.py b/tools/testing/source_code_test.py index 47c4c73c4d..5a546a9f48 100644 --- a/tools/testing/source_code_test.py +++ b/tools/testing/source_code_test.py @@ -19,6 +19,7 @@ from typedapi import ensure_api_is_typed import tensorflow_addons as tfa +import tensorflow as tf BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) @@ -39,6 +40,7 @@ def test_api_typed(): # Files within this list will be exempt from verification. exception_list = [ tfa.rnn.PeepholeLSTMCell, + tf.keras.optimizers.legacy.optimizer.Optimizer, ] help_message = ( "You can also take a look at the section about it in the CONTRIBUTING.md:\n" From 268d37aba881d13abd95b6b7b510e6cac3c02b3d Mon Sep 17 00:00:00 2001 From: chenmoneygithub Date: Sat, 18 Jun 2022 17:15:44 -0700 Subject: [PATCH 14/14] fix import --- tools/testing/source_code_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/testing/source_code_test.py b/tools/testing/source_code_test.py index 5a546a9f48..421c83d958 100644 --- a/tools/testing/source_code_test.py +++ b/tools/testing/source_code_test.py @@ -18,6 +18,7 @@ from typedapi import ensure_api_is_typed +import importlib import tensorflow_addons as tfa import tensorflow as tf @@ -40,8 +41,10 @@ def test_api_typed(): # Files within this list will be exempt from verification. exception_list = [ tfa.rnn.PeepholeLSTMCell, - tf.keras.optimizers.legacy.optimizer.Optimizer, ] + if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None: + exception_list.append(tf.keras.optimizers.legacy.Optimizer) + help_message = ( "You can also take a look at the section about it in the CONTRIBUTING.md:\n" "https://github.com/tensorflow/addons/blob/master/CONTRIBUTING.md#about-type-hints"