Repository: Lightning-AI/pytorch-lightning

Issue:
trainer.scale_batch_size() throws exception due to LRScheduler 
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

I tried finding the biggest possible batch_size for my training, but PL raises a `MisconfigurationException` saying that my `LRScheduler` (`ReduceLROnPlateau`) is conditioned on a metric that is only available after `validation_epoch_end`. The available metrics are: `loss, val_loss`.

I assume the LRScheduler requires a metric from the training loop for this to work? Why is this neccessary?
<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior:

1. Have a model with a metric that only exists in `validation_epoch_end`
2. Have a LRScheduler which monitors that metric
3. Use `trainer.scale_batch_size`
4. See error

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
```
File "C:\ProgramData\Anaconda3\envs\ml\lib\site-packages\pytorch_lightning\trainer\training_loop.py", line 779, in update_learning_rates
    raise MisconfigurationException(
pytorch_lightning.utilities.exceptions.MisconfigurationException: ReduceLROnPlateau conditioned on metric meanIoU which is not available. Available metrics are: loss,train_loss. Condition can be set using `monitor` key in lr scheduler dict

```

#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->
```
trainer = pl.Trainer(gpus=hparams.gpus)
new_batch_size = trainer.scale_batch_size(net, mode='binsearch', init_val=8)
```
and in my model:
```
    def configure_optimizers(self):
        opt = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        scheduler = {
         'scheduler': optim.lr_scheduler.ReduceLROnPlateau(opt, mode="max", factor=0.5, patience=5),
         'monitor': 'meanIoU',  # Default: val_loss
        }
        return [opt], [scheduler]

    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        iou_class, mean_iou = self.iou_metric.value()
        mean_iou = torch.tensor(mean_iou)
        self.iou_metric.reset()
        logs = {"val_loss": avg_loss, "meanIoU": mean_iou}
        return {"meanIoU": mean_iou, "log": logs,
                "progress_bar": {"val_loss": avg_loss, "meanIoU": mean_iou}}
```
### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
No Exception and the maximum batch_size for my model.

### Environment

* CUDA:
	- GPU:
		- GeForce RTX 2070 SUPER
	- available:         True
	- version:           10.1
* Packages:
	- numpy:             1.18.1
	- pyTorch_debug:     False
	- pyTorch_version:   1.4.0
	- pytorch-lightning: 0.7.6
	- tensorboard:       2.1.0
	- tqdm:              4.45.0
* System:
	- OS:                Windows
	- architecture:
		- 64bit
		- WindowsPE
	- processor:         AMD64 Family 23 Model 113 Stepping 0, AuthenticAMD
	- python:            3.8.2
	- version:           10.0.18362


### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Just to be sure, this error only happens when you run the `.scale_batch_size(...)` method and not when you run `.fit(...)`?
> Just to be sure, this error only happens when you run the `.scale_batch_size(...)` method and not when you run `.fit(...)`?

Correct
The problem here is that learning rates are updated right after the `training_epoch` has returned, even when we have reached `max_steps` and should just run the `training_teardown`
https://github.com/PyTorchLightning/pytorch-lightning/blob/ac76dfcf62a672c84f843f2e3158e4c6262776da/pytorch_lightning/trainer/training_loop.py#L348-L355
We could probably just switch order of the two statements and this will be solved.
Great that you found the reason for that behavior already!
@HansBambel can you locally try if this works for you, and if it does, would you be up for doing a PR?
> @HansBambel can you locally try if this works for you, and if it does, would you be up for doing a PR?

I'll try it out!

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -347,13 +347,13 @@ def train(self):
                 # -----------------
                 self.run_training_epoch()
 
-                # update LR schedulers
-                self.update_learning_rates(interval='epoch')
-
                 if self.max_steps and self.max_steps == self.global_step:
                     self.run_training_teardown()
                     return
 
+                # update LR schedulers
+                self.update_learning_rates(interval='epoch')
+
                 # early stopping
                 met_min_epochs = epoch >= self.min_epochs - 1
                 met_min_steps = self.global_step >= self.min_steps if self.min_steps else True


Test patch: