Preferred way to freeze a backbone, train for several steps, unfreeze the backbone, and continue training on all parameters.

[isaacgerg]

I see a lot of forum posts about freezing a backbone, training for several steps, then unfreezing the backbone, then continjue training. However, many posts seem to be missing steps (I believe) to account for batchnorm, dropout, etc. I am also confused about .freeze() and .eval() -- what are the differences and when should you use each?

To make this discussion concrete, below is an example pseudocode with my current understanding how to do this. In the example below, I train with a frozen backbone for 10 epochs and then train the entire network for the remaining 90 epochs . Is the pseudocode
correct?

class UnfreezeModelCallback(pl.Callback):
    def __init__(self, unfreeze_epoch: int):
        """
        Args:
            unfreeze_epoch (int): The epoch after which to unfreeze the model.
        """
        super().__init__()
        self.unfreeze_epoch = unfreeze_epoch

    def on_train_epoch_end(self, trainer, pl_module):
        # Check if the current epoch is the one to unfreeze the model
        if trainer.current_epoch == self.unfreeze_epoch:
            pl_module.backbone.unfreeze()

model = MyModel()  # LightningModule
model.backbone.freeze()  # Does this unfreeze recursively? i.e., LightningModules inside of LightningModules.  Why not just call eval() here? So many questions.

unfreeze_callback = UnfreezeModelCallback(10)

trainer = pl.Trainer(callbacks=[unfreeze_callback], max_epochs = 100)
trainer.fit(model, datamodule=...)