[Cleanup] change create_optimizer interface to be optimization config and runtime.

PiperOrigin-RevId: 351524712

official/core/base_task.py

@@ -24,7 +24,7 @@ from official.core import config_definitions
 from official.modeling import optimization
 from official.modeling import performance
 
-TrainerConfig = config_definitions.TrainerConfig
+OptimizationConfig = optimization.OptimizationConfig
 RuntimeConfig = config_definitions.RuntimeConfig
 
 
@@ -62,18 +62,18 @@ class Task(tf.Module, metaclass=abc.ABCMeta):
     return self._logging_dir
 
   @classmethod
-  def create_optimizer(cls, trainer_config: TrainerConfig,
+  def create_optimizer(cls, optimizer_config: OptimizationConfig,
                        runtime_config: Optional[RuntimeConfig] = None):
     """Creates an TF optimizer from configurations.
 
     Args:
-      trainer_config: the parameters of the trainer.
+      optimizer_config: the parameters of the Optimization settings.
       runtime_config: the parameters of the runtime.
 
     Returns:
       A tf.optimizers.Optimizer object.
     """
-    opt_factory = optimization.OptimizerFactory(trainer_config.optimizer_config)
+    opt_factory = optimization.OptimizerFactory(optimizer_config)
     optimizer = opt_factory.build_optimizer(opt_factory.build_learning_rate())
     # Configuring optimizer when loss_scale is set in runtime config. This helps
     # avoiding overflow/underflow for float16 computations.

official/core/base_trainer_test.py

@@ -59,7 +59,8 @@ class TrainerTest(tf.test.TestCase, parameterized.TestCase):
         config,
         task,
         model=task.build_model(),
-        optimizer=task.create_optimizer(config.trainer, config.runtime),
+        optimizer=task.create_optimizer(config.trainer.optimizer_config,
+                                        config.runtime),
         checkpoint_exporter=ckpt_exporter)
     return trainer
 
@@ -189,15 +190,18 @@ class TrainerTest(tf.test.TestCase, parameterized.TestCase):
                 }
             })))
     task = mock_task.MockTask(config.task, logging_dir=model_dir)
+
     def build_losses(labels, model_outputs, aux_losses=None):
       del labels, model_outputs
       return tf.constant([np.nan], tf.float32) + aux_losses
+
     task.build_losses = build_losses
     trainer = trainer_lib.Trainer(
         config,
         task,
         model=task.build_model(),
-        optimizer=task.create_optimizer(config.trainer, config.runtime))
+        optimizer=task.create_optimizer(config.trainer.optimizer_config,
+                                        config.runtime))
     trainer.add_recovery(config.trainer, checkpoint_manager=checkpoint_manager)
     with self.assertRaises(RuntimeError):
       _ = trainer.train(tf.convert_to_tensor(2, dtype=tf.int32))

official/core/train_utils.py

@@ -134,7 +134,8 @@ def create_trainer(params: config_definitions.ExperimentConfig,
   """Create trainer."""
   logging.info('Running default trainer.')
   model = task.build_model()
-  optimizer = task.create_optimizer(params.trainer, params.runtime)
+  optimizer = task.create_optimizer(params.trainer.optimizer_config,
+                                    params.runtime)
   return trainer_cls(
       params,
       task,

official/modeling/multitask/multitask.py

@@ -21,10 +21,9 @@ from official.core import base_task
 from official.core import config_definitions
 from official.core import task_factory
 from official.modeling import optimization
-from official.modeling import performance
 from official.modeling.multitask import configs
 
-TrainerConfig = config_definitions.TrainerConfig
+OptimizationConfig = optimization.OptimizationConfig
 RuntimeConfig = config_definitions.RuntimeConfig
 
 
@@ -105,28 +104,11 @@ class MultiTask(tf.Module, metaclass=abc.ABCMeta):
     return self._task_weights[task_name]
 
   @classmethod
-  def create_optimizer(cls, trainer_config: TrainerConfig,
+  def create_optimizer(cls,
+                       optimizer_config: OptimizationConfig,
                        runtime_config: Optional[RuntimeConfig] = None):
-    """Creates an TF optimizer from configurations.
-
-    Args:
-      trainer_config: the parameters of the trainer.
-      runtime_config: the parameters of the runtime.
-
-    Returns:
-      A tf.optimizers.Optimizer object.
-    """
-    opt_factory = optimization.OptimizerFactory(trainer_config.optimizer_config)
-    optimizer = opt_factory.build_optimizer(opt_factory.build_learning_rate())
-    # Configuring optimizer when loss_scale is set in runtime config. This helps
-    # avoiding overflow/underflow for float16 computations.
-    if runtime_config and runtime_config.loss_scale:
-      optimizer = performance.configure_optimizer(
-          optimizer,
-          use_float16=runtime_config.mixed_precision_dtype == "float16",
-          loss_scale=runtime_config.loss_scale)
-
-    return optimizer
+    return base_task.Task.create_optimizer(
+        optimizer_config=optimizer_config, runtime_config=runtime_config)
 
   def joint_train_step(self, task_inputs, multi_task_model, optimizer,
                        task_metrics):

official/modeling/multitask/train_lib.py

@@ -49,7 +49,8 @@ def run_experiment_wtih_multitask_eval(
   is_training = 'train' in mode
   is_eval = 'eval' in mode
   with distribution_strategy.scope():
-    optimizer = train_task.create_optimizer(params.trainer, params.runtime)
+    optimizer = train_task.create_optimizer(params.trainer.optimizer_config,
+                                            params.runtime)
     model = train_task.build_model()
     if is_training:
       trainer = core_lib.Trainer(