[WIP] Gradual Unfreezing to mitigate catastrophic forgetting

ludwig/modules/gradual_unfreezer.py

@@ -0,0 +1,28 @@
+from ludwig.schema.gradual_unfreezer import GradualUnfreezerConfig
+
+
+class GradualUnfreezer:
+ def __init__(self, config: GradualUnfreezerConfig, model):
+ self.config = config
+ self.model = model
+ self.thaw_epochs = self.config.thaw_epochs
+ self.layers_to_thaw = self.config.layers_to_thaw
+
+ if len(self.thaw_epochs) != len(self.layers_to_thaw):
+ raise ValueError("The length of thaw_epochs and layers_to_thaw must be equal.")
+ self.layers = dict(zip(self.thaw_epochs, self.layers_to_thaw))
+
+ def thaw(self, current_epoch: int) -> None:
+ if current_epoch in self.layers:
+ current_layers = self.layers[current_epoch]
+ for layer in current_layers:
+ self.thawParameter(layer)
+
+ # thaw individual layer
+ def thawParameter(self, layer):
+ # is there a better way to do this instead of iterating through all parameters?
+ for name, p in self.model.named_parameters():
+ if layer in str(name):
+ p.requires_grad_(True)
+ else:
+ raise ValueError("Layer type doesn't exist within model architecture")

ludwig/schema/gradual_unfreezer.py

@@ -0,0 +1,78 @@
+from abc import ABC
+from dataclasses import field
+from typing import Dict
+
+from marshmallow import fields, ValidationError
+
+import ludwig.schema.utils as schema_utils
+from ludwig.api_annotations import DeveloperAPI
+from ludwig.constants import MODEL_ECD
+from ludwig.schema.metadata import TRAINER_METADATA
+from ludwig.schema.utils import ludwig_dataclass
+
+
+@DeveloperAPI
+@ludwig_dataclass
+class GradualUnfreezerConfig(schema_utils.BaseMarshmallowConfig, ABC):
+ """Configuration for gradual unfreezing parameters."""
+
+ thaw_epochs: list = schema_utils.List(
+ int,
+ default=None,
+ description="Epochs to thaw at. For example, [1, 2, 3, 4] will thaw layers in layers_to_thaw 2D array",
+ parameter_metadata=TRAINER_METADATA[MODEL_ECD]["gradual_unfreezer"]["thaw_epochs"],
+ )
+
+ layers_to_thaw: list = schema_utils.List(
+ list,
+ inner_type=str,
+ default=None,
+ description="Individual layers to thaw at each thaw_epoch. 2D Array",
+ parameter_metadata=TRAINER_METADATA[MODEL_ECD]["gradual_unfreezer"]["layers_to_thaw"],
+ )
+
+
+@DeveloperAPI
+def GradualUnfreezerDataclassField(description: str, default: Dict = None):
+ allow_none = True
+ default = default or {}
+
+ class GradualUnfreezerMarshmallowField(fields.Field):
+ def _deserialize(self, value, attr, data, **kwargs):
+ if value is None:
+ return value
+ if isinstance(value, dict):
+ try:
+ return GradualUnfreezerConfig.Schema().load(value)
+ except (TypeError, ValidationError) as e:
+ raise ValidationError(
+ f"Invalid params for gradual unfreezer: {value}, see GradualUnfreezerConfig class. Error: {e}"
+ )
+ raise ValidationError("Field should be None or dict")
+
+ def _jsonschema_type_mapping(self):
+ return {
+ **schema_utils.unload_jsonschema_from_marshmallow_class(GradualUnfreezerConfig),
+ "title": "gradual_unfreeze_options",
+ "description": description,
+ }
+
+ if not isinstance(default, dict):
+ raise ValidationError(f"Invalid default: `{default}`")
+
+ load_default = lambda: GradualUnfreezerConfig.Schema().load(default)
+ dump_default = GradualUnfreezerConfig.Schema().dump(default)
+
+ return field(
+ metadata={
+ "marshmallow_field": GradualUnfreezerMarshmallowField(
+ allow_none=allow_none,
+ load_default=load_default,
+ dump_default=dump_default,
+ metadata={
+ "description": description,
+ },
+ )
+ },
+ default_factory=load_default,
+ )

ludwig/schema/metadata/configs/trainer.yaml

@@ -640,6 +640,11 @@ ecd:
  eta_min:
  expected_impact: 1
  ui_display_name: Eta Min
+ gradual_unfreezer:
+ thaw_epochs:
+ expected_impact: 1
+ layers_to_thaw:
+ expected_impact: 1
 gbm:
  learning_rate:
  commonly_used: true

ludwig/schema/trainer.py

@@ -18,6 +18,7 @@
 )
 from ludwig.error import ConfigValidationError
 from ludwig.schema import utils as schema_utils
+from ludwig.schema.gradual_unfreezer import GradualUnfreezerConfig, GradualUnfreezerDataclassField
 from ludwig.schema.lr_scheduler import LRSchedulerConfig, LRSchedulerDataclassField
 from ludwig.schema.metadata import TRAINER_METADATA
 from ludwig.schema.optimizers import (
@@ -177,6 +178,11 @@ def __post_init__(self):
  ],
  )
 
+ gradual_unfreezer: GradualUnfreezerConfig = GradualUnfreezerDataclassField(
+ description="Parameter values for gradual unfreezer.",
+ default=None,
+ )
+
  learning_rate_scheduler: LRSchedulerConfig = LRSchedulerDataclassField(
  description="Parameter values for learning rate scheduler.",
  default=None,

ludwig/trainers/trainer.py

@@ -57,6 +57,7 @@
 from ludwig.models.ecd import ECD
 from ludwig.models.llm import LLM
 from ludwig.models.predictor import Predictor
+from ludwig.modules.gradual_unfreezer import GradualUnfreezer
 from ludwig.modules.lr_scheduler import LRScheduler
 from ludwig.modules.metric_modules import get_improved_fn, get_initial_validation_value
 from ludwig.modules.metric_registry import get_metric_objective
@@ -215,6 +216,7 @@ def __init__(
  self.dist_model = None
  self.optimizer = None
  self.scheduler = None
+ self.gradual_unfreezer = None
 
  self.prepare()
 
@@ -1002,6 +1004,11 @@ def train(
  total_steps=self.total_steps,
  )
 
+ # Initialize gradual unfreezer
+ if self.config.gradual_unfreezer.thaw_epochs:
+ self.gradual_unfreezer = GradualUnfreezer(self.config.gradual_unfreezer, self.model)
+ logger.info(f"Gradual unfreezing for {len(self.gradual_unfreezer.thaw_epochs)} epoch(s)")
+
  if self.is_coordinator():
  logger.info(
  f"Training for {self.total_steps} step(s), approximately "
@@ -1029,7 +1036,12 @@ def train(
  if profiler:
  profiler.start()
 
+ current_epoch = 0
+
  while progress_tracker.steps < self.total_steps:
+ if self.gradual_unfreezer:
+ self.gradual_unfreezer.thaw(current_epoch)
+
  # note that batch size may change over epochs
  batcher.set_epoch(progress_tracker.epoch, progress_tracker.batch_size)
 
@@ -1086,6 +1098,8 @@ def train(
  # Early stop if needed.
  if should_break:
  break
+
+ current_epoch += 1
  finally:
  # ================ Finished Training ================
  self.callback(

tests/ludwig/config_sampling/static_schema.json

@@ -145285,6 +145285,83 @@
  "title": "learning_rate_scheduler_options",
  "type": "object"
  },
+ "gradual_unfreezer": {
+ "additionalProperties": true,
+ "description": "Parameter values for gradual unfreezing.",
+ "properties": {
+ "thaw_epochs": {
+ "default": null,
+ "description": "List of epochs to unfreeze layers",
+ "items": {
+ "title": "thaw_epochs",
+ "type": "integer"
+ },
+ "parameter_metadata": {
+ "commonly_used": false,
+ "compute_tier": 0,
+ "default_value_reasoning": null,
+ "description_implications": null,
+ "example_value": null,
+ "expected_impact": 3,
+ "internal_only": false,
+ "literature_references": [
+ "https://aclanthology.org/P18-1031.pdf",
+ "https://arxiv.org/pdf/2301.05487.pdf"
+ ],
+ "long_description": "",
+ "other_information": null,
+ "related_parameters": null,
+ "short_description": "",
+ "suggested_values": null,
+ "suggested_values_reasoning": null,
+ "ui_display_name": "Thaw Epochs"
+ },
+ "title": "thaw_epochs",
+ "type": [
+ "integer",
+ "null"
+ ]
+ },
+ "layers_to_thaw": {
+ "default": [],
+ "description": "List of layers to thaw at each epoch",
+ "items": {
+ "type": "array",
+ "items": {
+ "type": "string",
+ "title": "layer_name"
+ }
+ },
+ "parameter_metadata": {
+ "commonly_used": false,
+ "compute_tier": 0,
+ "default_value_reasoning": null,
+ "description_implications": null,
+ "example_value": null,
+ "expected_impact": 3,
+ "internal_only": false,
+ "literature_references": [
+ "https://aclanthology.org/P18-1031.pdf",
+ "https://arxiv.org/pdf/2301.05487.pdf"
+ ],
+ "long_description": "",
+ "other_information": null,
+ "related_parameters": null,
+ "short_description": "",
+ "suggested_values": null,
+ "suggested_values_reasoning": null,
+ "ui_display_name": "Layers To Thaw"
+ },
+ "title": "layers_to_thaw",
+ "type": [
+ "array",
+ "null"
+ ]
+ }
+ },
+ "title": "gradual_unfreeze_options",
+ "type": "object"
+ },
  "max_batch_size": {
  "default": 1099511627776,
  "description": "Auto batch size tuning and increasing batch size on plateau will be capped at this value. The default value is 2^40.",

tests/ludwig/modules/test_gradual_unfreezing.py

@@ -0,0 +1,27 @@
+from ludwig.encoders.image.torchvision import TVSwinTransformerEncoder
+from ludwig.modules.gradual_unfreezer import GradualUnfreezer, GradualUnfreezerConfig
+from ludwig.utils.misc_utils import set_random_seed
+
+
+def test_gradual_unfreezer():
+ set_random_seed(13)
+
+ model = TVSwinTransformerEncoder(
+ model_variant="t",
+ use_pretrained=False,
+ saved_weights_in_checkpoint=True,
+ trainable=False,
+ )
+ config = GradualUnfreezerConfig(thaw_epochs=[1, 2], layers_to_thaw=[["features.0", "features.1"], ["features.2"]])
+
+ unfreezer = GradualUnfreezer(config=config, model=model)
+
+ for epoch in range(10):
+ unfreezer.thaw(epoch)
+
+ for name, p in model.named_parameters():
+ layer_to_thaw = any(layer in str(name) for layer_list in config.layers_to_thaw for layer in layer_list)
+ if layer_to_thaw:
+ assert p.requires_grad
+ else:
+ assert not p.requires_grad