Repository: Lightning-AI/pytorch-lightning

Issue:
Add strict option to scheduler
## 🚀 Feature

When a `ReduceLROnPlateau` scheduler is used, training stops with a raised `MisconfigurationException` if the `monitor` is not found. We should add a `strict` option (same as in the `EarlyStopping` callback) so it is configurable.

EarlyStopping:
https://github.com/PyTorchLightning/pytorch-lightning/blob/c64520e658806a87f282c74299b2ea4b7ea493ea/pytorch_lightning/callbacks/early_stopping.py#L61-L62
https://github.com/PyTorchLightning/pytorch-lightning/blob/c64520e658806a87f282c74299b2ea4b7ea493ea/pytorch_lightning/callbacks/early_stopping.py#L107-L116

OptimizerConnector:
https://github.com/PyTorchLightning/pytorch-lightning/blob/c64520e658806a87f282c74299b2ea4b7ea493ea/pytorch_lightning/trainer/connectors/optimizer_connector.py#L53-L59

### Motivation

If my first epoch fails to produce a valid loss, the scheduler stops training.
### Pitch

```python
    def configure_optimizers(self):
        optimizer = ...
        scheduler = {
            "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
                optimizer,
                mode="min",
            ),
            "monitor": "checkpoint_on",
            "interval": "epoch",
            "frequency": 1,
            "strict": False  # Allow this
        }
        return [optimizer], [scheduler]
```




Discussion:
> If my first epoch fails to produce a valid loss, the scheduler stops training.

can you elaborate a bit more on this case, not sure what is the case that you would endup without loss...
@carmocca is there a specific reason why your model fails to produce a validation loss for the first epoch?
I like this addition as we have already encountered before that some users setup their training such that sometimes the validation loss is computed after they need it for their scheduler (https://github.com/PyTorchLightning/pytorch-lightning/issues/1598, https://github.com/PyTorchLightning/pytorch-lightning/issues/1156)
Same reason as https://github.com/PyTorchLightning/pytorch-lightning/pull/3527#issuecomment-694827383

> To give a concrete example for why I need this:
>
> I am training a model using variable length data (text) with the CTC loss.
> The CTC loss produces infinite losses when the model output is too short to be aligned to the targets.
> If that happens for a full batch, then I have no loss.
> And you cant check at the dataloader level if a batch produced long enough model outputs.

Resolving patch:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -972,7 +972,8 @@ def configure_optimizers(
                     'interval': 'epoch', # The unit of the scheduler's step size
                     'frequency': 1, # The frequency of the scheduler
                     'reduce_on_plateau': False, # For ReduceLROnPlateau scheduler
-                    'monitor': 'val_loss' # Metric to monitor
+                    'monitor': 'val_loss', # Metric for ReduceLROnPlateau to monitor
+                    'strict': True # Whether to crash the training if `monitor` is not found
                 }
 
             If user only provides LR schedulers, then their configuration will set to default as shown above.
diff --git a/pytorch_lightning/trainer/connectors/optimizer_connector.py b/pytorch_lightning/trainer/connectors/optimizer_connector.py
--- a/pytorch_lightning/trainer/connectors/optimizer_connector.py
+++ b/pytorch_lightning/trainer/connectors/optimizer_connector.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
@@ -57,12 +57,20 @@ def update_learning_rates(self, interval: str, monitor_metrics=None):
                         monitor_val = self.trainer.logger_connector.callback_metrics.get(monitor_key)
 
                     if monitor_val is None:
-                        avail_metrics = ','.join(list(self.trainer.logger_connector.callback_metrics.keys()))
-                        raise MisconfigurationException(
+                        if lr_scheduler.get('strict', True):
+                            avail_metrics = self.trainer.logger_connector.callback_metrics.keys()
+                            raise MisconfigurationException(
+                                f'ReduceLROnPlateau conditioned on metric {monitor_key}'
+                                f' which is not available. Available metrics are: {avail_metrics}.'
+                                ' Condition can be set using `monitor` key in lr scheduler dict'
+                            )
+                        rank_zero_warn(
                             f'ReduceLROnPlateau conditioned on metric {monitor_key}'
-                            f' which is not available. Available metrics are: [{avail_metrics}].'
-                            ' Condition can be set using `monitor` key in lr scheduler dict'
+                            ' which is not available but strict is set to `False`.'
+                            ' Skipping learning rate update.',
+                            RuntimeWarning,
                         )
+                        continue
                     # update LR
                     old_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr']
                     lr_scheduler['scheduler'].step(monitor_val)


Test patch:
diff --git a/tests/trainer/test_optimizers.py b/tests/trainer/test_optimizers.py
--- a/tests/trainer/test_optimizers.py
+++ b/tests/trainer/test_optimizers.py
@@ -360,3 +360,36 @@ def configure_optimizers(self):
         weights_summary=None,
     )
     trainer.fit(model)
+
+
+def test_lr_scheduler_strict(tmpdir):
+    """
+    Test "strict" support in lr_scheduler dict
+    """
+    model = EvalModelTemplate()
+    optimizer = torch.optim.Adam(model.parameters())
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
+
+    model.configure_optimizers = lambda: {
+        'optimizer': optimizer,
+        'lr_scheduler': {'scheduler': scheduler, 'monitor': 'giraffe', 'strict': True},
+    }
+    with pytest.raises(
+        MisconfigurationException,
+        match=r'ReduceLROnPlateau conditioned on metric .* which is not available\. Available metrics are:',
+    ):
+        trainer.fit(model)
+
+    model.configure_optimizers = lambda: {
+        'optimizer': optimizer,
+        'lr_scheduler': {
+            'scheduler': scheduler,
+            'monitor': 'giraffe',
+            'strict': False,
+        },
+    }
+    with pytest.warns(
+        RuntimeWarning, match=r'ReduceLROnPlateau conditioned on metric .* which is not available but strict'
+    ):
+        assert trainer.fit(model)