Internal change

PiperOrigin-RevId: 319058156

Author: Allen Wang

official/vision/image_classification/classifier_trainer.py

@@ -339,7 +339,8 @@ def train_and_eval(
     optimizer = optimizer_factory.build_optimizer(
         optimizer_name=params.model.optimizer.name,
         base_learning_rate=learning_rate,
-        params=params.model.optimizer.as_dict())
+        params=params.model.optimizer.as_dict(),
+        model=model)
 
     metrics_map = _get_metrics(one_hot)
     metrics = [metrics_map[metric] for metric in params.train.metrics]

official/vision/image_classification/optimizer_factory.py

@@ -18,11 +18,12 @@
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
+from typing import Any, Dict, Text, List
+
 from absl import logging
 import tensorflow as tf
 import tensorflow_addons as tfa
 
-from typing import Any, Dict, Text, List
 from official.vision.image_classification import learning_rate
 from official.vision.image_classification.configs import base_configs
 
@@ -250,7 +251,8 @@ def from_config(cls, config, custom_objects=None):
 def build_optimizer(
     optimizer_name: Text,
     base_learning_rate: tf.keras.optimizers.schedules.LearningRateSchedule,
-    params: Dict[Text, Any]):
+    params: Dict[Text, Any],
+    model: tf.keras.Model = None):
   """Build the optimizer based on name.
 
   Args:
@@ -261,6 +263,8 @@ def build_optimizer(
     params: String -> Any dictionary representing the optimizer params.
       This should contain optimizer specific parameters such as
       `base_learning_rate`, `decay`, etc.
+    model: The `tf.keras.Model`. This is used for the shadow copy if using
+      `MovingAverage`.
 
   Returns:
     A tf.keras.Optimizer.
@@ -322,10 +326,13 @@ def build_optimizer(
   # Moving average should be applied last, as it's applied at test time
   moving_average_decay = params.get('moving_average_decay', 0.)
   if moving_average_decay is not None and moving_average_decay > 0.:
+    if model is None:
+      raise ValueError('`model` must be provided if using `MovingAverage`.')
     logging.info('Including moving average decay.')
     optimizer = MovingAverage(
-        optimizer,
+        optimizer=optimizer,
         average_decay=moving_average_decay)
+    optimizer.shadow_copy(model)
   return optimizer
 
 

official/vision/image_classification/optimizer_factory_test.py

@@ -19,15 +19,21 @@
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
-import tensorflow as tf
-
 from absl.testing import parameterized
+
+import tensorflow as tf
 from official.vision.image_classification import optimizer_factory
 from official.vision.image_classification.configs import base_configs
 
 
 class OptimizerFactoryTest(tf.test.TestCase, parameterized.TestCase):
 
+  def build_toy_model(self) -> tf.keras.Model:
+    """Creates a toy `tf.Keras.Model`."""
+    model = tf.keras.Sequential()
+    model.add(tf.keras.layers.Dense(1, input_shape=(1,)))
+    return model
+
   @parameterized.named_parameters(
       ('sgd', 'sgd', 0., False),
       ('momentum', 'momentum', 0., False),
@@ -40,6 +46,7 @@ class OptimizerFactoryTest(tf.test.TestCase, parameterized.TestCase):
       ('rmsprop_ema', 'rmsprop', 0.999, False))
   def test_optimizer(self, optimizer_name, moving_average_decay, lookahead):
     """Smoke test to be sure no syntax errors."""
+    model = self.build_toy_model()
     params = {
         'learning_rate': 0.001,
         'rho': 0.09,
@@ -51,7 +58,8 @@ def test_optimizer(self, optimizer_name, moving_average_decay, lookahead):
     optimizer = optimizer_factory.build_optimizer(
         optimizer_name=optimizer_name,
         base_learning_rate=params['learning_rate'],
-        params=params)
+        params=params,
+        model=model)
     self.assertTrue(issubclass(type(optimizer), tf.keras.optimizers.Optimizer))
 
   def test_unknown_optimizer(self):