Repository: sql-machine-learning/elasticdl

Issue:
Master extracts optimizer type and related hyper-parameters (learning rate etc) from python code and sends them to Go PS as arguments



Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/model_utils.py b/elasticdl/python/common/model_utils.py
--- a/elasticdl/python/common/model_utils.py
+++ b/elasticdl/python/common/model_utils.py
@@ -1,6 +1,8 @@
 import importlib.util
 import os
 
+import tensorflow as tf
+
 from elasticdl.python.common.log_utils import default_logger as logger
 from elasticdl.python.data.odps_io import is_odps_configured
 from elasticdl.python.worker.prediction_outputs_processor import (
@@ -179,3 +181,30 @@ def get_non_embedding_trainable_vars(model, embedding_layers):
         if not is_embedding_item:
             non_embedding_trainable_vars.append(var)
     return non_embedding_trainable_vars
+
+
+def get_optimizer_info(optimizer):
+    """
+    Get optimizer type and its argument values
+    """
+    # TODO: add more supported optimizers
+    OPT_ARGUMENTS = {
+        "SGD": ["learning_rate", "momentum", "nesterov"],
+        "Adam": ["learning_rate", "beta_1", "beta_2", "epsilon", "amsgrad"],
+        "unkown": [],
+    }
+    opt_type = "unknown"
+    opt_argument = ""
+
+    if isinstance(optimizer, tf.keras.optimizers.SGD):
+        opt_type = "SGD"
+    elif isinstance(optimizer, tf.keras.optimizers.Adam):
+        opt_type = "Adam"
+    opt_config = optimizer.get_config()
+    for arg_name in OPT_ARGUMENTS[opt_type]:
+        arg_value = opt_config[arg_name]
+        # For callable, only get the value from 1st call
+        if callable(arg_value):
+            arg_value = arg_value()
+        opt_argument += arg_name + "=" + str(arg_value) + ";"
+    return opt_type, opt_argument
diff --git a/elasticdl/python/master/master.py b/elasticdl/python/master/master.py
--- a/elasticdl/python/master/master.py
+++ b/elasticdl/python/master/master.py
@@ -21,6 +21,7 @@
 from elasticdl.python.common.model_utils import (
     get_dict_from_params_str,
     get_module_file_path,
+    get_optimizer_info,
     load_model_from_module,
     load_module,
 )
@@ -340,6 +341,7 @@ def _create_instance_manager(self, args):
             worker_args.extend(build_arguments_from_parsed_result(args))
 
             if args.use_go_ps:
+                opt_type, opt_args = get_optimizer_info(self.optimizer)
                 # TODO: rename the Go PS executable using a meaningful filename
                 ps_command = ["main"]
                 ps_args = [
@@ -361,6 +363,8 @@ def _create_instance_manager(self, args):
                     "-keep_checkpoint_max=" + str(args.keep_checkpoint_max),
                     "-checkpoint_dir_for_init="
                     + str(args.checkpoint_dir_for_init),
+                    "-opt_type=" + opt_type,
+                    "-opt_args=" + opt_args,
                 ]
             else:
                 ps_command = ["python"]


Test patch:
diff --git a/elasticdl/python/tests/model_utils_test.py b/elasticdl/python/tests/model_utils_test.py
--- a/elasticdl/python/tests/model_utils_test.py
+++ b/elasticdl/python/tests/model_utils_test.py
@@ -1,11 +1,14 @@
 import os
 import unittest
 
+import tensorflow as tf
+
 from elasticdl.python.common.model_utils import (
     _get_spec_value,
     get_dict_from_params_str,
     get_model_spec,
     get_module_file_path,
+    get_optimizer_info,
 )
 
 _model_zoo_path = os.path.dirname(os.path.realpath(__file__))
@@ -102,6 +105,50 @@ def test_get_dict_from_params_str(self):
         )
         self.assertEqual(get_dict_from_params_str(""), {})
 
+    def test_get_optimizer_info(self):
+        learning_rate = 0.1
+        momentum = 0.0
+        nesterov = False
+        expected_args = (
+            "learning_rate="
+            + str(learning_rate)
+            + ";momentum="
+            + str(momentum)
+            + ";nesterov=False;"
+        )
+        opt = tf.keras.optimizers.SGD(
+            learning_rate=learning_rate, momentum=momentum, nesterov=nesterov
+        )
+        opt_type, opt_args = get_optimizer_info(opt)
+        self.assertEqual(opt_type, "SGD")
+        self.assertEqual(opt_args, expected_args)
+
+        beta_1 = 0.8
+        beta_2 = 0.6
+        epsilon = 1e-08
+        amsgrad = False
+        expected_args = (
+            "learning_rate="
+            + str(learning_rate)
+            + ";beta_1="
+            + str(beta_1)
+            + ";beta_2="
+            + str(beta_2)
+            + ";epsilon="
+            + str(epsilon)
+            + ";amsgrad=False;"
+        )
+        opt = tf.keras.optimizers.Adam(
+            learning_rate=learning_rate,
+            beta_1=beta_1,
+            beta_2=beta_2,
+            epsilon=epsilon,
+            amsgrad=amsgrad,
+        )
+        opt_type, opt_args = get_optimizer_info(opt)
+        self.assertEqual(opt_type, "Adam")
+        self.assertEqual(opt_args, expected_args)
+
 
 if __name__ == "__main__":
     unittest.main()