Freeze backbone support for RetinaNet

PiperOrigin-RevId: 443733011

official/vision/configs/retinanet.py

@@ -167,6 +167,10 @@ class RetinaNetTask(cfg.TaskConfig):
   # If set, the Waymo Open Dataset evaluator would be used.
   use_wod_metrics: bool = False
 
+  # If set, freezes the backbone during training.
+  # TODO(crisnv) Add paper link when available.
+  freeze_backbone: bool = False
+
 
 @exp_factory.register_config_factory('retinanet')
 def retinanet() -> cfg.ExperimentConfig:

official/vision/tasks/retinanet.py

@@ -59,6 +59,10 @@ class RetinaNetTask(base_task.Task):
         input_specs=input_specs,
         model_config=self.task_config.model,
         l2_regularizer=l2_regularizer)
+
+    if self.task_config.freeze_backbone:
+      model.backbone.trainable = False
+
     return model
 
   def initialize(self, model: tf.keras.Model):