Fix BaseFinetuning re-freezing unfrozen modules on resume

[yentur]

What does this PR do?

Fixes #21901

BaseFinetuning restores half of its state when a run is resumed. setup() calls freeze_before_training() on every fit, including a resumed one, because it has to run before configure_optimizers(). on_fit_start() then puts the saved param groups back on the optimizer, but it only restores group membership, never param.requires_grad.

The result is that modules unfrozen before the interruption come back frozen. They sit in the optimizer with the right learning rate and get no gradients, so they never update again. BackboneFinetuning.finetune_function() unfreezes only on the exact epoch epoch == unfreeze_backbone_at_epoch and only rescales the learning rate afterwards, so nothing puts requires_grad back.

This restores requires_grad for the param groups that were added after the ones configure_optimizers() created. Those are precisely the groups unfreeze_and_add_param_group() appended in the previous run, so the frozen state is restored as it was rather than blanket-unfrozen. Resuming before the unfreeze epoch adds no extra groups and therefore changes nothing.

Same 4 epochs, tracking requires_grad and whether the backbone weights actually change, interrupted after epoch 1:

                uninterrupted           resumed at epoch 2
 epoch  requires_grad  weights_moved requires_grad  weights_moved
     0          False          False         False          False
     1           True           True          True           True
     2           True           True         False          False     <- before
     3           True           True         False          False     <- before

After the fix the resumed columns match the uninterrupted ones for every epoch.

With a schedule that unfreezes block i at epoch i, resuming at epoch 2:

uninterrupted : [True, True, True, False]
resumed before: [False, False, True, False]
resumed after : [True, True, True, False]

Known limitation, not addressed here: freeze_module() also sets track_running_stats = False on BatchNorm layers, and that flag is not recoverable from the saved param groups. A BatchNorm that was unfrozen with train_bn=False will have its affine parameters trainable again but keep track_running_stats=False, and one with affine=False has no parameters in any group at all. Restoring that would need the callback to persist which modules it unfroze, which is a larger change to the checkpoint contents.

The bug only appears when the unfrozen modules were added to the optimizer as a new param group, which is the pattern the BackboneFinetuning docstring shows. If configure_optimizers() already hands every parameter to the optimizer, unfreeze_and_add_param_group() drops them in filter_on_optimizer(), no group is recorded, and there is nothing in the metadata to restore from.

test_callbacks_restore_backbone already walks this resume path but asserts nothing after the second fit(), which is why the suite stayed green. The new test covers both directions: unfrozen stays unfrozen, and still-frozen stays frozen.

$ pytest callbacks/test_finetuning_callback.py -q
10 passed, 1 skipped, 34 warnings in 1.10s

With only the change to finetuning.py reverted:

$ pytest callbacks/test_finetuning_callback.py::test_finetuning_restores_requires_grad_on_resume -q
FAILED callbacks/test_finetuning_callback.py::test_finetuning_restores_requires_grad_on_resume[1-True] - assert False is True
1 failed, 1 passed

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.