Transfer learning
The problem
The standard flow for transfer learning is as follows:

Start from a pretrained model and create a new head
Freeze all layers but the head
Train
Unfreeze all (or some) layers
Train (with differential learning rates)

Let's call each "freeze/unfreeze -> train" step a phase.
What can change between phases?

The lr_scheduler should restart
New parameter groups can be added/removed to the optimizer and to the lr_scheduler (We don't really add a parameter group to the lr_scheduler, but we should pass a list of learning rates, one for each group (this is called differential learning rates))

Parameter groups
I've implemented a first idea on how to implement this as a ModuleMixin here, the idea is for the user to override model_splits, which returns the parameter groups of the model.
You can see how this gets used for differential learning rates here

Changing phases
There are two main ways of changing phases, but currently none of them are smooth.
Use multiple trainers
Use a separate trainer for each phase, always resuming the last state by using resume_from_checkpoint.
A scratch notebook for this idea is here
The problem is that trainer also resumes the lr_scheduler state, but we need it to start from zero
Implement the logic in on_epoch_start
Idea notebook

Check the current_epoch and change phases when needed
Might need to add/remove parameter groups from the optimizer/lr_scheduler
Changing the scheduler is causing problems when used together with LeraningRateLogger