Change object detector learning rate decay to cosine decay.

PiperOrigin-RevId: 527337105

mediapipe/model_maker/python/vision/object_detector/hyperparameters.py

@@ -14,7 +14,7 @@
 """Hyperparameters for training object detection models."""
 
 import dataclasses
-from typing import List
+from typing import Optional
 
 from mediapipe.model_maker.python.core import hyperparameters as hp
 
@@ -29,12 +29,13 @@ class HParams(hp.BaseHParams):
     epochs: Number of training iterations over the dataset.
     do_fine_tuning: If true, the base module is trained together with the
       classification layer on top.
-    learning_rate_epoch_boundaries: List of epoch boundaries where
-      learning_rate_epoch_boundaries[i] is the epoch where the learning rate
-      will decay to learning_rate * learning_rate_decay_multipliers[i].
-    learning_rate_decay_multipliers: List of learning rate multipliers which
-      calculates the learning rate at the ith boundary as learning_rate *
-      learning_rate_decay_multipliers[i].
+    cosine_decay_epochs: The number of epochs for cosine decay learning rate.
+      See
+      https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/CosineDecay
+        for more info.
+    cosine_decay_alpha: The alpha value for cosine decay learning rate. See
+      https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/CosineDecay
+        for more info.
   """
 
   # Parameters from BaseHParams class.
@@ -42,41 +43,9 @@ class HParams(hp.BaseHParams):
   batch_size: int = 32
   epochs: int = 10
 
-  # Parameters for learning rate decay
-  learning_rate_epoch_boundaries: List[int] = dataclasses.field(
-      default_factory=lambda: []
-  )
-  learning_rate_decay_multipliers: List[float] = dataclasses.field(
-      default_factory=lambda: []
-  )
-
-  def __post_init__(self):
-    # Validate stepwise learning rate parameters
-    lr_boundary_len = len(self.learning_rate_epoch_boundaries)
-    lr_decay_multipliers_len = len(self.learning_rate_decay_multipliers)
-    if lr_boundary_len != lr_decay_multipliers_len:
-      raise ValueError(
-          "Length of learning_rate_epoch_boundaries and ",
-          "learning_rate_decay_multipliers do not match: ",
-          f"{lr_boundary_len}!={lr_decay_multipliers_len}",
-      )
-    # Validate learning_rate_epoch_boundaries
-    if (
-        sorted(self.learning_rate_epoch_boundaries)
-        != self.learning_rate_epoch_boundaries
-    ):
-      raise ValueError(
-          "learning_rate_epoch_boundaries is not in ascending order: ",
-          self.learning_rate_epoch_boundaries,
-      )
-    if (
-        self.learning_rate_epoch_boundaries
-        and self.learning_rate_epoch_boundaries[-1] > self.epochs
-    ):
-      raise ValueError(
-          "Values in learning_rate_epoch_boundaries cannot be greater ",
-          "than epochs",
-      )
+  # Parameters for cosine learning rate decay
+  cosine_decay_epochs: Optional[int] = None
+  cosine_decay_alpha: float = 0.0
 
 
 @dataclasses.dataclass

mediapipe/model_maker/python/vision/object_detector/object_detector.py

@@ -354,19 +354,16 @@ class ObjectDetector(classifier.Classifier):
       A tf.keras.optimizer.Optimizer for model training.
     """
     init_lr = self._hparams.learning_rate * self._hparams.batch_size / 256
-    if self._hparams.learning_rate_epoch_boundaries:
-      lr_values = [init_lr] + [
-          init_lr * m for m in self._hparams.learning_rate_decay_multipliers
-      ]
-      lr_step_boundaries = [
-          steps_per_epoch * epoch_boundary
-          for epoch_boundary in self._hparams.learning_rate_epoch_boundaries
-      ]
-      learning_rate = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
-          lr_step_boundaries, lr_values
-      )
-    else:
-      learning_rate = init_lr
+    decay_epochs = (
+        self._hparams.cosine_decay_epochs
+        if self._hparams.cosine_decay_epochs
+        else self._hparams.epochs
+    )
+    learning_rate = tf.keras.optimizers.schedules.CosineDecay(
+        init_lr,
+        steps_per_epoch * decay_epochs,
+        self._hparams.cosine_decay_alpha,
+    )
     return tf.keras.optimizers.experimental.SGD(
         learning_rate=learning_rate, momentum=0.9
     )