Repository: tensorflow/addons

Issue:
Requesting support for decaying the weight decay constant in weight_decay_optimizers.
API type: Optimizer
Who will benefit: Anyone using a weight decay optimizer who wants to use learning rate decay.

The comments for the weight decay optimizers state that if learning rate decay is employed the weight decay scaling parameter should also be decayed with the same function (albeit most probably different initial and maximal values of the decay parameter.  With tfa 0.7.0, learning rate schedule  is supported, which it wasn't in 0.6.0, but it isn't for the weight decay parameter. I think that there needs to be a test for the weight decay being an instance of learning rate schedule class and executing the __call__ function with the global step when that is the case.



Discussion:
cc @PhilJd 

Resolving patch:
diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers.py b/tensorflow_addons/optimizers/weight_decay_optimizers.py
--- a/tensorflow_addons/optimizers/weight_decay_optimizers.py
+++ b/tensorflow_addons/optimizers/weight_decay_optimizers.py
@@ -75,8 +75,9 @@ def __init__(self, weight_decay: Union[FloatTensorLike, Callable], **kwargs):
         """Extension class that adds weight decay to an optimizer.
 
         Args:
-            weight_decay: A `Tensor` or a floating point value, the factor by
-                which a variable is decayed in the update step.
+            weight_decay: A `Tensor`, a floating point value, or a schedule
+                that is a `tf.keras.optimizers.schedules.LearningRateSchedule`
+                to decay the variable by, in the update step.
             **kwargs: Optional list or tuple or set of `Variable` objects to
                 decay.
         """
@@ -92,6 +93,23 @@ def get_config(self):
         )
         return config
 
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        # LR handling copied from optimizer_v2.OptimizerV2
+        if "learning_rate" in config:
+            if isinstance(config["learning_rate"], dict):
+                config["learning_rate"] = tf.keras.optimizers.schedules.deserialize(
+                    config["learning_rate"], custom_objects=custom_objects
+                )
+
+        if "weight_decay" in config:
+            if isinstance(config["weight_decay"], dict):
+                config["weight_decay"] = tf.keras.optimizers.schedules.deserialize(
+                    config["weight_decay"], custom_objects=custom_objects
+                )
+
+        return cls(**config)
+
     def minimize(self, loss, var_list, grad_loss=None, name=None, decay_var_list=None):
         """Minimize `loss` by updating `var_list`.
 
@@ -151,19 +169,21 @@ def apply_gradients(self, grads_and_vars, name=None, decay_var_list=None, **kwar
 
     def _decay_weights_op(self, var):
         if not self._decay_var_list or var.ref() in self._decay_var_list:
-            return var.assign_sub(
-                self._get_hyper("weight_decay", var.dtype) * var, self._use_locking
-            )
+            return var.assign_sub(self._decayed_wd(var.dtype) * var, self._use_locking)
         return tf.no_op()
 
     def _decay_weights_sparse_op(self, var, indices):
         if not self._decay_var_list or var.ref() in self._decay_var_list:
-            update = -self._get_hyper("weight_decay", var.dtype) * tf.gather(
-                var, indices
-            )
+            update = -self._decayed_wd(var.dtype) * tf.gather(var, indices)
             return self._resource_scatter_add(var, indices, update)
         return tf.no_op()
 
+    def _decayed_wd(self, var_dtype):
+        wd_t = self._get_hyper("weight_decay", var_dtype)
+        if isinstance(wd_t, tf.keras.optimizers.schedules.LearningRateSchedule):
+            wd_t = tf.cast(wd_t(self.iterations), var_dtype)
+        return wd_t
+
     # Here, we overwrite the apply functions that the base optimizer calls.
     # super().apply_x resolves to the apply_x function of the BaseOptimizer.
 


Test patch:
diff --git a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
--- a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
+++ b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
@@ -251,6 +251,21 @@ def test_keras_fit():
     model.fit(x, y, epochs=1)
 
 
+def test_keras_fit_with_schedule():
+    """Check if calling model.fit works with wd schedule."""
+    model = tf.keras.models.Sequential([tf.keras.layers.Dense(2)])
+    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+    wd_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
+        WEIGHT_DECAY, decay_steps=10, decay_rate=0.9
+    )
+    optimizer = weight_decay_optimizers.AdamW(
+        learning_rate=1e-4, weight_decay=wd_schedule
+    )
+    model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])
+    x, y = np.random.uniform(size=(2, 4, 1))
+    model.fit(x, y, epochs=1)
+
+
 @pytest.mark.parametrize("dtype", [(tf.half, 0), (tf.float32, 1), (tf.float64, 2)])
 def test_sparse_sgdw(dtype):
     do_test(
@@ -355,3 +370,15 @@ def test_serialization():
     config = tf.keras.optimizers.serialize(optimizer)
     new_optimizer = tf.keras.optimizers.deserialize(config)
     assert new_optimizer.get_config() == optimizer.get_config()
+
+
+def test_serialization_with_wd_schedule():
+    wd_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
+        WEIGHT_DECAY, decay_steps=10, decay_rate=0.9
+    )
+    optimizer = weight_decay_optimizers.AdamW(
+        learning_rate=1e-4, weight_decay=wd_schedule
+    )
+    config = tf.keras.optimizers.serialize(optimizer)
+    new_optimizer = tf.keras.optimizers.deserialize(config)
+    assert new_optimizer.get_config() == optimizer.get_config()