Fix learning rate scheduler for the CNN model of cifar10 (sql-machine…

…-learning#2296)

* Fix learning rate scheduler for the cifar10 model

* Add an unittest for error

elasticdl/python/elasticdl/callbacks.py

@@ -16,7 +16,6 @@
 
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras import backend as K
 
 from elasticdl.proto import elasticdl_pb2
 from elasticdl.python.common.constants import Mode
@@ -151,4 +150,4 @@ def on_train_batch_begin(self, batch, logs=None):
             raise ValueError(
                 'The output of the "schedule" function should be float.'
             )
-        K.set_value(self.model.optimizer.lr, K.get_value(lr))
+        self.model.optimizer.lr = lr

elasticdl/python/tests/callbacks_test.py

@@ -118,6 +118,20 @@ class LearningRateSchedulerTest(unittest.TestCase):
     def _schedule(self, model_version):
         return 0.2 if model_version < 2 else 0.1
 
+    def test_raise_error(self):
+        def _schedule(model_version):
+            return 1 if model_version < 2 else 2
+
+        learning_rate_scheduler = LearningRateScheduler(_schedule)
+        model = tf.keras.Model()
+        learning_rate_scheduler.set_model(model)
+        with self.assertRaises(ValueError):
+            learning_rate_scheduler.on_train_batch_begin(batch=1)
+
+        model.optimizer = tf.optimizers.SGD(0.1)
+        with self.assertRaises(ValueError):
+            learning_rate_scheduler.on_train_batch_begin(batch=1)
+
     def test_learning_rate_scheduler(self):
         learning_rate_scheduler = LearningRateScheduler(self._schedule)
         model = tf.keras.Model()

model_zoo/cifar10/cifar10_functional_api.py

@@ -138,7 +138,7 @@ def _schedule(model_version):
         else:
             return 0.001
 
-    LearningRateScheduler(_schedule)
+    return [LearningRateScheduler(_schedule)]
 
 
 def dataset_fn(dataset, mode, _):

model_zoo/cifar10/cifar10_resnet50.py

@@ -53,7 +53,7 @@ def _schedule(model_version):
         else:
             return 0.001
 
-    LearningRateScheduler(_schedule)
+    return [LearningRateScheduler(_schedule)]
 
 
 def dataset_fn(dataset, mode, _):