Remember the eval mode of submodules when switching trainer stages

[awaelchli]

What does this PR do?

Fixes #18930
Part of #16827

A common issue users are facing is that the loop calls train() on the LightningModule despite the user having frozen certain layers. For example,

• user finetunes a model and freezes certain layers they don't want to train
• user's model has a feature extractor they don't train

This leads to a surprise when the user finds out that their batch norm layers have changed statistics, even though they were set explicitly to eval() mode. To avoid this, the user has to learn that they should override the on_validation_model_eval() and on_validation_model_train() hooks in the module, but this is a detail difficult to find in our docs and get right. Most users who face this challenge end up on slack or GH to ask for help.

The PR makes the following changes to automate this for the user:

• The validation loop captures the .training mode of every submodule before calling .eval() now. When the validation loop ends, and before switching to training, it restores the .training mode on all submodules to what it was before. This ensures that layers the user has chosen to be in eval mode remain in eval mode!
• The fit-loop no longer calls .train() at the beginning with the same motivation: The user can now set a subset of their model to .eval() mode / freeze it explicitly in the LightningModule's __init__ without doing acrobatics with hooks, and the Trainer will respect it and preserve it (see the added test). Note: This is not a breaking change, because PyTorch's default is to have a model in .training=True mode.

---

📚 Documentation preview 📚: https://pytorch-lightning--18951.org.readthedocs.build/en/18951/

cc @Borda @justusschock @awaelchli

[codecov]

Codecov Report

Merging #18951 (e71ab68) into master (b80107e) will decrease coverage by 27%.
Report is 1 commits behind head on master.
The diff coverage is 100%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #18951      +/-   ##
==========================================
- Coverage      76%      48%     -27%     
==========================================
  Files         450      442       -8     
  Lines       36508    36383     -125     
==========================================
- Hits        27583    17572   -10011     
- Misses       8925    18811    +9886     

[carmocca]

Does this PR replace #18826?

[awaelchli]

Does this PR replace #18826?

@carmocca Good find! Yes in fact it does, I verified thanks to the repro example that was posted there. After this PR lands, we could still use #18826 to add an additional test case for the tuner.

[adosar]

@awaelchli

This leads to a surprise when the user finds out that their batch norm layers have changed statistics, even though they were set explicitly to eval() mode. To avoid this, the user has to learn that they should override the on_validation_model_eval() and on_validation_model_train()

Since the problem was in the training loop which called train(), even if the user sets them in on_validation_model_{eval|train}, it would lead to incorrect behavior (before this PR). Am I right? (just asking for better understanding Lightning)

Also, currently in the docs it is shown that model.train() is called after the validation ends what was done in this PR.

# ...
for batch_idx, batch in enumerate(train_dataloader):
    loss = model.training_step(batch, batch_idx)
    loss.backward()
    # ...

    if validate_at_some_point:
        # disable grads + batchnorm + dropout
        torch.set_grad_enabled(False)
        model.eval()

        # ----------------- VAL LOOP ---------------
        for val_batch_idx, val_batch in enumerate(val_dataloader):
            val_out = model.validation_step(val_batch, val_batch_idx)
        # ----------------- VAL LOOP ---------------

        # enable grads + batchnorm + dropout
        torch.set_grad_enabled(True)
        model.train()

I would happy to open a new issue for fixing the docs to better align with this very useful PR!