Internal change

PiperOrigin-RevId: 466508813

official/modeling/optimization/ema_optimizer.py

@@ -21,7 +21,7 @@ import tensorflow as tf
 # pylint: disable=protected-access
 
 
-class ExponentialMovingAverage(tf.keras.optimizers.Optimizer):
+class ExponentialMovingAverage(tf.keras.optimizers.legacy.Optimizer):
   """Optimizer that computes an exponential moving average of the variables.
 
   Empirically it has been found that using the moving average of the trained

official/modeling/optimization/lars_optimizer.py

@@ -22,7 +22,7 @@ import tensorflow as tf
 # pylint: disable=protected-access
 
 
-class LARS(tf.keras.optimizers.Optimizer):
+class LARS(tf.keras.optimizers.legacy.Optimizer):
   """Layer-wise Adaptive Rate Scaling for large batch training.
 
   Introduced by "Large Batch Training of Convolutional Networks" by Y. You,

official/modeling/optimization/legacy_adamw.py

@@ -20,7 +20,7 @@ from absl import logging
 import tensorflow as tf
 
 
-class AdamWeightDecay(tf.keras.optimizers.Adam):
+class AdamWeightDecay(tf.keras.optimizers.legacy.Adam):
   """Adam enables L2 weight decay and clip_by_global_norm on gradients.
 
   [Warning!]: Keras optimizer supports gradient clipping and has an AdamW

official/vision/beta/projects/yolo/optimization/sgd_torch.py

@@ -43,7 +43,7 @@ def _var_key(var):
   return var._unique_id
 
 
-class SGDTorch(tf.keras.optimizers.Optimizer):
+class SGDTorch(tf.keras.optimizers.legacy.Optimizer):
   """Optimizer that simulates the SGD module used in pytorch.