Repository: Lightning-AI/pytorch-lightning

Issue:
Trainer(precision=16) fails with optim.lr_scheduler.ReduceLROnPlateau
<!-- 
### 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

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior:

1. Create a `pl.LightningModule` that returns your optimizer along with a `optim.lr_scheduler.ReduceLROnPlateau` scheduler from `configure_optimizers`
2. Create a `pl.Trainer` wit `precision=16`
3. Run your training (i.e., `trainer.fit(model)`)
4. See error

```console
Traceback (most recent call last):                                                                                                  
  File "main.py", line 65, in <module>                                                                                              
    main()                                                                                                                          
  File "main.py", line 61, in main                                                                                                  
    trainer.fit(model)                                                                                                              
  File "/workspace/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 889, in fit                             
    self.dp_train(model)                                                                                                            
  File "/workspace/pytorch-lightning/pytorch_lightning/trainer/distrib_parts.py", line 223, in dp_train                  
    self.reinit_scheduler_properties(optimizers, self.lr_schedulers)                                                                
  File "/workspace/pytorch-lightning/pytorch_lightning/trainer/optimizers.py", line 122, in reinit_scheduler_properties  
    scheduler.__class__.__mro__[idx].__init__(scheduler, optimizer)                                                                 
UnboundLocalError: local variable 'idx' referenced before assignment                                                                
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


<!-- #### Code sample -->
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->

<!-- ### Expected behavior -->

<!-- A clear and concise description of what you expected to happen. -->

<!-- ### Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py
```
 - PyTorch Version (1.5):
 - OS (Linux):

### Additional context
-->

<!-- Add any other context about the problem here. -->

The error occurs in `pytorch-lightning/pytorch_lightning/trainer/optimizers.py", line 122`.

```python
def reinit_scheduler_properties(self, optimizers: list, schedulers: list):
    # Reinitialize optimizer.step properties added by schedulers
    for scheduler in schedulers:
        for optimizer in optimizers:
            scheduler = scheduler['scheduler']
            # check that we dont mix users optimizers and schedulers
            if scheduler.optimizer == optimizer:
                # Find the mro belonging to the base lr scheduler class
                for i, mro in enumerate(scheduler.__class__.__mro__):
                    if mro == optim.lr_scheduler._LRScheduler:
                        idx = i
                scheduler.__class__.__mro__[idx].__init__(scheduler, optimizer)
```
The `idx` local variable is unassigned because `optim.lr_scheduler.ReduceLROnPlateau` is not a subclass of `optim.lr_scheduler._LRScheduler`.

I could work around the error by adding a specific check for `optim.lr_scheduler.ReduceLROnPlateau` but I'm not sure if this is a good solution.

```python
def reinit_scheduler_properties(self, optimizers: list, schedulers: list):
    # Reinitialize optimizer.step properties added by schedulers
    for scheduler in schedulers:
        for optimizer in optimizers:
            scheduler = scheduler['scheduler']
            # check that we dont mix users optimizers and schedulers
            if scheduler.optimizer == optimizer:
                # Find the mro belonging to the base lr scheduler class
                for i, mro in enumerate(scheduler.__class__.__mro__):
                    if mro == optim.lr_scheduler._LRScheduler:
                        idx = i
                    elif mro == optim.lr_scheduler.ReduceLROnPlateau:
                        idx = i
                scheduler.__class__.__mro__[idx].__init__(scheduler, optimizer)
```

### Related issue in PyTorch:
ReduceLROnPlateau parent class is not _LRScheduler #21981
https://github.com/pytorch/pytorch/issues/21981


Discussion:
Hi! thanks for your contribution!, great first issue!
@naokishibuya good catch. It seems like a problem that should be solved upstream in pytorch, but for now we can solve this locally. Would you be up for a PR?
When I tried this fix, it solved the error but unfortunately `ReduceLROnPlateau` stopped working for me (i.e. there was no indication of the LR decreasing with `verbose=True` or on TensorBoard). If I switched back to `precision=32`, it works normally again
I think that the fix is actually working, however only calling `__init__(scheduler, optimizer)` will reset all other arguments (patience, mode, ect) to default values for the `ReduceLrOnPlauteau` scheduler. A solution to this is to copy over these properties:
```
__init__(scheduler, optimizer, patience=scheduler.patience,mode=scheduler.mode,...)
```
Again I think this is a bit hacky, and a proper solution upstream in pytorch is better. 
I think this does the trick for me:
```python
def reinit_scheduler_properties(self, optimizers: list, schedulers: list):
    # Reinitialize optimizer.step properties added by schedulers
    for scheduler in schedulers:
        for optimizer in optimizers:
            scheduler = scheduler["scheduler"]
            # check that we dont mix users optimizers and schedulers
            if scheduler.optimizer == optimizer:
                # Find the mro belonging to the base lr scheduler class
                for i, mro in enumerate(scheduler.__class__.__mro__):
                    if (
                        mro == optim.lr_scheduler._LRScheduler
                        or mro == optim.lr_scheduler.ReduceLROnPlateau
                    ):
                        idx = i
                        state = scheduler.state_dict()
                    else:
                        state = None
                scheduler.__class__.__mro__[idx].__init__(scheduler, optimizer)
                if state is not None:
                    scheduler.load_state_dict(state)
```
Happy to open a PR if it looks ok to you guys

Resolving patch:
diff --git a/pytorch_lightning/trainer/optimizers.py b/pytorch_lightning/trainer/optimizers.py
--- a/pytorch_lightning/trainer/optimizers.py
+++ b/pytorch_lightning/trainer/optimizers.py
@@ -111,15 +111,25 @@ def configure_schedulers(self, schedulers: list):
     def reinit_scheduler_properties(self, optimizers: list, schedulers: list):
         # Reinitialize optimizer.step properties added by schedulers
         for scheduler in schedulers:
+            scheduler = scheduler['scheduler']
+
             for optimizer in optimizers:
-                scheduler = scheduler['scheduler']
                 # check that we dont mix users optimizers and schedulers
                 if scheduler.optimizer == optimizer:
                     # Find the mro belonging to the base lr scheduler class
                     for i, mro in enumerate(scheduler.__class__.__mro__):
-                        if mro == optim.lr_scheduler._LRScheduler:
+                        if (
+                            mro == optim.lr_scheduler._LRScheduler
+                            or mro == optim.lr_scheduler.ReduceLROnPlateau
+                        ):
                             idx = i
-                    scheduler.__class__.__mro__[idx].__init__(scheduler, optimizer)
+                            state = scheduler.state_dict()
+                        else:
+                            state = None
+
+                scheduler.__class__.__mro__[idx].__init__(scheduler, optimizer)
+                if state is not None:
+                    scheduler.load_state_dict(state)
 
 
 class _MockOptimizer(Optimizer):


Test patch: