Move some gradient ascent optimizer flags into the gin config

This patch moves command line flags in es_trainer_lib to set the gradient
ascent optimizer and gradient ascent optimizer flags to the gin config along
with all the other model/optimizer hyperparameters. These don't make a lot of
sense to have as command line parameters given their nature and everything
else similar being defined in gin configs.

Reviewers: mtrofin

Reviewed By: mtrofin

Pull Request: #433

Author: boomanaiden154

compiler_opt/es/es_trainer_lib.py

@@ -14,6 +14,7 @@
 """Local ES trainer."""
 
 from absl import flags, logging
+import enum
 import functools
 import gin
 import tensorflow as tf
@@ -31,22 +32,12 @@
 
 FLAGS = flags.FLAGS
 
-_BETA1 = flags.DEFINE_float("beta1", 0.9,
-                            "Beta1 for ADAM gradient ascent optimizer.")
-_BETA2 = flags.DEFINE_float("beta2", 0.999,
-                            "Beta2 for ADAM gradient ascent optimizer.")
 _GRAD_REG_ALPHA = flags.DEFINE_float(
     "grad_reg_alpha", 0.01,
     "Weight of regularization term in regression gradient.")
 _GRAD_REG_TYPE = flags.DEFINE_string(
     "grad_reg_type", "ridge",
     "Regularization method to use with regression gradient.")
-_GRADIENT_ASCENT_OPTIMIZER_TYPE = flags.DEFINE_string(
-    "gradient_ascent_optimizer_type", None,
-    "Gradient ascent optimization algorithm: 'momentum' or 'adam'")
-flags.mark_flag_as_required("gradient_ascent_optimizer_type")
-_MOMENTUM = flags.DEFINE_float(
-    "momentum", 0.0, "Momentum for momentum gradient ascent optimizer.")
 _OUTPUT_PATH = flags.DEFINE_string("output_path", "",
                                    "Path to write all output")
 _PRETRAINED_POLICY_PATH = flags.DEFINE_string(
@@ -60,11 +51,22 @@
                                      "List of paths to training corpora")
 
 
+@gin.constants_from_enum(module="es_trainer_lib")
+class GradientAscentOptimizerType(enum.Enum):
+  INVALID = 0
+  MOMENTUM = enum.auto()
+  ADAM = enum.auto()
+
+
 @gin.configurable
 def train(additional_compilation_flags=(),
           delete_compilation_flags=(),
           replace_compilation_flags=(),
-          worker_class=None):
+          worker_class=None,
+          beta1=0.9,
+          beta2=0.999,
+          momentum=0.0,
+          gradient_ascent_optimizer_type=GradientAscentOptimizerType.ADAM):
   """Train with ES."""
 
   if not _TRAIN_CORPORA.value:
@@ -130,21 +132,20 @@ def train(additional_compilation_flags=(),
   # TODO(linzinan): delete all unused parameters.
 
   # ------------------ GRADIENT ASCENT OPTIMIZERS ------------------------------
-  if _GRADIENT_ASCENT_OPTIMIZER_TYPE.value == "momentum":
+  if gradient_ascent_optimizer_type == GradientAscentOptimizerType.MOMENTUM:
     logging.info("Running momentum gradient ascent optimizer")
     # You can obtain a vanilla gradient ascent optimizer by setting momentum=0.0
     # and setting step_size to the desired learning rate.
     gradient_ascent_optimizer = (
         gradient_ascent_optimization_algorithms.MomentumOptimizer(
-            learner_config.step_size, _MOMENTUM.value))
-  elif _GRADIENT_ASCENT_OPTIMIZER_TYPE.value == "adam":
+            learner_config.step_size, momentum))
+  elif gradient_ascent_optimizer_type == GradientAscentOptimizerType.ADAM:
     logging.info("Running Adam gradient ascent optimizer")
     gradient_ascent_optimizer = (
         gradient_ascent_optimization_algorithms.AdamOptimizer(
-            learner_config.step_size, _BETA1.value, _BETA2.value))
+            learner_config.step_size, beta1, beta2))
   else:
-    logging.info("No gradient ascent \
-                 optimizer selected. Stopping.")
+    logging.info("No gradient ascent optimizer selected. Stopping.")
     return
   # ----------------------------------------------------------------------------