Add freeze_backbone flag for image_classification.

PiperOrigin-RevId: 460570094

official/projects/vit/configs/image_classification.py

@@ -63,6 +63,7 @@ class ImageClassificationTask(cfg.TaskConfig):
   evaluation: Evaluation = Evaluation()
   init_checkpoint: Optional[str] = None
   init_checkpoint_modules: str = 'all'  # all or backbone
+  freeze_backbone: bool = False
 
 
 IMAGENET_TRAIN_EXAMPLES = 1281167

official/vision/configs/image_classification.py

@@ -97,6 +97,7 @@ class ImageClassificationTask(cfg.TaskConfig):
   init_checkpoint_modules: str = 'all'  # all or backbone
   model_output_keys: Optional[List[int]] = dataclasses.field(
       default_factory=list)
+  freeze_backbone: bool = False
 
 
 @exp_factory.register_config_factory('image_classification')

official/vision/tasks/image_classification.py

@@ -49,6 +49,9 @@ class ImageClassificationTask(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):