Fine-tuning support

realfake/callbacks.py

@@ -47,6 +47,12 @@ class FeatureExtractorFreezeUnfreeze(BaseFinetuning):
     def freeze_before_training(self, pl_module: pl.LightningModule) -> None:
         rank_zero_info("Freezing backbone")
         self.freeze(_get_backbone(pl_module.model))
+        enabled_layers = [
+            name
+            for name, child in pl_module.model.named_children()
+            if all(param.requires_grad for param in child.parameters())
+        ]
+        rank_zero_info(f"Gradient enabled layers: [{', '.join(enabled_layers)}]")
 
     def finetune_function(self, pl_module: "pl.LightningModule", epoch: int, optimizer, opt_idx: int) -> None:
         if epoch == self._unfreeze_at_epoch:

realfake/models.py

@@ -35,6 +35,7 @@ class RealFakeParams(Args):
     epochs: int = Field(6)
     base_lr: float = Field(1e-3)
     pretrained: bool = Field(True)
+    progress_bar: bool = Field(False)
     accelerator: AcceleratorParams = Field(default_factory=AcceleratorParams)
 
 

realfake/train.py

@@ -7,7 +7,7 @@ import pytorch_lightning as pl
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.plugins.environments import SLURMEnvironment
 
-from realfake.callbacks import ConsoleLogger
+from realfake.callbacks import ConsoleLogger, FeatureExtractorFreezeUnfreeze
 from realfake.models import RealFakeParams
 from realfake.utils import get_checkpoints_dir, find_latest_checkpoint
 
@@ -45,7 +45,7 @@ def prepare_trainer(args: RealFakeParams) -> pl.Trainer:
         max_epochs=args.epochs,
         num_nodes=1,
         num_sanity_val_steps=0,
-        enable_progress_bar=False,
+        enable_progress_bar=args.progress_bar,
         callbacks=[
             ConsoleLogger(),
             ModelCheckpoint(
@@ -56,12 +56,13 @@ def prepare_trainer(args: RealFakeParams) -> pl.Trainer:
                 dirpath=checkpoints_dir,
                 filename="%s-{epoch:02d}-{val_acc:.4f}" % args.model_name,
             ),
+            FeatureExtractorFreezeUnfreeze(unfreeze_at_epoch=args.freeze_epochs)
         ],
         resume_from_checkpoint=existing_checkpoint,
     )
 
     if job_id is not None:
-        trainer_params["enable_progress_bar"] = True
+        trainer_params["enable_progress_bar"] = False
         trainer_params["plugins"] = SLURMEnvironment(requeue_signal=signal.SIGHUP)
         trainer_params["strategy"] = args.accelerator.strategy