Repository: keras-team/keras (Python)
Date: 2024-04-10 22:45:57+00:00
Task type: Bug Fix

Issue:
keras adamw optimizer failed with callable parameters in TensorFlow2.16
When we were working on upgrading keras 2 to keras 3 in TensorFlow plugin, one of our adamw related unit test failed, which is a sub unit test using callable  lambda as learning_rate argument. We also found this ut failed in TensorFlow2.16 official docker image. The error log is :
![image](https://github.com/keras-team/keras/assets/25453568/3f212c38-f1d1-4dc0-95bd-f8f3d8c37916)


```python
"""Tests for adam optimizer with weight decay."""

import numpy as np
import keras
import tensorflow as tf
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import constant_op
from tensorflow.python.platform import test
from tensorflow.python.framework import test_util

from keras.src.optimizers import adamw

DATA_TYPES = [
    dtypes.float32
]

WEIGHT_DECAY = 0.1

def adamw_update_numpy(
    param, grad_t, slot_vars, learning_rate, beta_1, beta_2, epsilon, weight_decay, amsgrad
):
    """Numpy update function for AdamW."""
    lr, beta1, beta2, eps, wd = (
        v() if callable(v) else v
        for v in (learning_rate, beta_1, beta_2, epsilon, weight_decay)
    )
    t = slot_vars.get("t", 0) + 1
    lr_t = lr * np.sqrt(1 - beta2 ** t) / (1 - beta1 ** t)
    slot_vars["m"] = beta1 * slot_vars.get("m", 0) + (1 - beta1) * grad_t
    slot_vars["v"] = beta2 * slot_vars.get("v", 0) + (1 - beta2) * grad_t ** 2
    if amsgrad:
        slot_vars["v_hat"] = slot_vars.get("v_hat", 0)
        slot_vars["v_hat"] = np.maximum(slot_vars["v_hat"], slot_vars["v"])
        param_t = param * (1 - wd * lr) - lr_t * slot_vars["m"] / (np.sqrt(slot_vars["v_hat"]) + eps)
    else:
        param_t = param * (1 - wd * lr) - lr_t * slot_vars["m"] / (np.sqrt(slot_vars["v"]) + eps)
    slot_vars["t"] = t
    return param_t, slot_vars

class AdamWeightDecayOptimizerTest(test_util.TensorFlowTestCase):

    def doTestBasic(self, use_callable_params=False, do_sparse=False, do_amsgrad=False):
        for dtype in DATA_TYPES:
            # Initialize variables for numpy implementation.
            np_slot_vars0, np_slot_vars1 = {}, {}
            var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
            grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
            var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
            grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)

            # Create Tensorflow variables.
            itex_var0 = tf.Variable(var0_np)
            itex_var1 = tf.Variable(var1_np)

            # Adapt callable parameters
            learning_rate = lambda: 0.01
            beta_1=lambda: 0.9
            beta_2=lambda: 0.999
            if not use_callable_params:
                learning_rate = learning_rate()
                beta_1 = beta_1()
                beta_2 = beta_2()

            # Adapt sparse
            if do_sparse:
                grads0_np_indices = np.array([0, 1], dtype=np.int32)
                grads0 = tf.IndexedSlices(
                    tf.constant(grads0_np), tf.constant(grads0_np_indices), tf.constant([2])
                )
                grads1_np_indices = np.array([0, 1], dtype=np.int32)
                grads1 = tf.IndexedSlices(
                    tf.constant(grads1_np), tf.constant(grads1_np_indices), tf.constant([2])
                )
            else:
                grads0 = constant_op.constant(grads0_np)
                grads1 = constant_op.constant(grads1_np)

            adamw_opt = adamw.AdamW(weight_decay=WEIGHT_DECAY, learning_rate=learning_rate, amsgrad=do_amsgrad)

            # Run 3 steps of the optimizer
            for _ in range(3):
                adamw_opt.apply_gradients(
                    zip([grads0, grads1], [itex_var0, itex_var1])
                )
                var0_np, np_slot_vars0 = adamw_update_numpy(
                    var0_np, grads0_np, np_slot_vars0, weight_decay=WEIGHT_DECAY, learning_rate=learning_rate, 
                    beta_1=beta_1, beta_2=beta_2, epsilon=1e-7, amsgrad=do_amsgrad)
                var1_np, np_slot_vars1 = adamw_update_numpy(
                    var1_np, grads1_np, np_slot_vars1, weight_decay=WEIGHT_DECAY, learning_rate=learning_rate, 
                    beta_1=beta_1, beta_2=beta_2, epsilon=1e-7, amsgrad=do_amsgrad)
                # Validate updated parameters
                self.assertAllCloseAccordingToType(itex_var0.numpy(), var0_np)
                self.assertAllCloseAccordingToType(itex_var1.numpy(), var1_np)


    def testCallableParamsAdamW(self):
        '''ResourceApplyAdamWithWeightDecay is a DPCPP op, don't have cpu registration 
            TODO: waiting for CPU registration of ResourceApplyAdamWithWeightDecay then enable
            this test case on CPU'''
        if not test.is_gpu_available():
            self.skipTest("No GPU available")
        self.doTestBasic(use_callable_params=True)


if __name__ == "__main__":
    test.main()
```



Issue discussion:
https://github.com/keras-team/keras/blob/6c591d7d34c3ffaa50e805fd75c83d9c2a23414f/keras/optimizers/base_optimizer.py#L560
Here is the root cause. If learning_rate is a callable object, then it doesn't need any arguments. 
I might give this one a stab if no one picks it up.
@kapoor1992 , You can create a PR
@sachinprasadhs Will do :) 

Solution patch:
diff --git a/keras/optimizers/base_optimizer.py b/keras/optimizers/base_optimizer.py
--- a/keras/optimizers/base_optimizer.py
+++ b/keras/optimizers/base_optimizer.py
@@ -567,7 +567,7 @@ def _get_current_learning_rate(self):
         ):
             return self._learning_rate(self.iterations)
         elif callable(self._learning_rate):
-            return self._learning_rate(self.iterations)
+            return self._learning_rate()
         return self._learning_rate
 
     def _filter_empty_gradients(self, grads, vars):


Test patch:
diff --git a/keras/optimizers/optimizer_test.py b/keras/optimizers/optimizer_test.py
--- a/keras/optimizers/optimizer_test.py
+++ b/keras/optimizers/optimizer_test.py
@@ -243,3 +243,12 @@ def test_tf_checkpointing(self):
         checkpoint.restore(save_path)
         pred = model.predict(x)
         self.assertAllClose(pred, ref_pred, atol=1e-5)
+
+    def test_callable_learning_rate(self):
+        v = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
+        grads = backend.convert_to_tensor([[1.0, 1.0], [1.0, 1.0]])
+        optimizer = optimizers.AdamW(learning_rate=lambda: 0.0001)
+        self.assertAllClose(optimizer.iterations, 0)
+        optimizer.apply_gradients([(grads, v)])
+        self.assertAllClose(v, [[1.0, 2.0], [3.0, 4.0]], atol=1e-4)
+        self.assertAllClose(optimizer.iterations, 1)