Bugfix/5487 auto lr ordering

pytorch_lightning/tuner/lr_finder.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import importlib
 import os
-from typing import List, Optional, Sequence, Union, Callable
+from typing import List, Optional, Sequence, Union, Callable, Any
 from functools import wraps
 
 import numpy as np
@@ -40,8 +40,28 @@
     from tqdm import tqdm
 
 
+def __choose_lr_assigner(trainer, model: LightningModule) -> Callable[[Any], None]:
+    if isinstance(trainer.auto_lr_find, str):
+        if not lightning_hasattr(model, trainer.auto_lr_find):
+            raise MisconfigurationException(
+                f'`auto_lr_find` was set to {trainer.auto_lr_find}, however'
+                ' could not find this as a field in `model` or `model.hparams`.')
+        return lambda val: lightning_setattr(model, trainer.auto_lr_find, val)
+
+    if lightning_hasattr(model, 'lr'):
+        return lambda val: lightning_setattr(model, 'lr', val)
+    if lightning_hasattr(model, 'learning_rate'):
+        return lambda val: lightning_setattr(model, 'learning_rate', val)
+    raise MisconfigurationException(
+        'When auto_lr_find is set to True, expects that `model` or'
+        ' `model.hparams` either has field `lr` or `learning_rate`'
+        ' that can overridden')
+
+
 def _run_lr_finder_internally(trainer, model: LightningModule):
     """ Call lr finder internally during Trainer.fit() """
+    lr_assigner = __choose_lr_assigner(trainer, model)
+
     lr_finder = lr_find(trainer, model)
 
     if lr_finder is None:
@@ -50,24 +70,8 @@ def _run_lr_finder_internally(trainer, model: LightningModule):
     lr = lr_finder.suggestion()
 
     # TODO: log lr.results to self.logger
-    if isinstance(trainer.auto_lr_find, str):
-        # Try to find requested field, may be nested
-        if lightning_hasattr(model, trainer.auto_lr_find):
-            lightning_setattr(model, trainer.auto_lr_find, lr)
-        else:
-            raise MisconfigurationException(
-                f'`auto_lr_find` was set to {trainer.auto_lr_find}, however'
-                ' could not find this as a field in `model` or `model.hparams`.')
-    else:
-        if lightning_hasattr(model, 'lr'):
-            lightning_setattr(model, 'lr', lr)
-        elif lightning_hasattr(model, 'learning_rate'):
-            lightning_setattr(model, 'learning_rate', lr)
-        else:
-            raise MisconfigurationException(
-                'When auto_lr_find is set to True, expects that `model` or'
-                ' `model.hparams` either has field `lr` or `learning_rate`'
-                ' that can overridden')
+    lr_assigner(lr)
+
     log.info(f'Learning rate set to {lr}')
 
 

tests/trainer/test_dataloaders.py

@@ -1148,7 +1148,6 @@ def test_replace_sampler_with_multiprocessing_context(tmpdir):
     train = DataLoader(train, batch_size=32, num_workers=2, multiprocessing_context=context, shuffle=True)
 
     class ExtendedBoringModel(BoringModel):
-
         def train_dataloader(self):
             return train
 

tests/trainer/test_lr_finder.py

@@ -18,8 +18,12 @@
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning import seed_everything
 from tests.base import EvalModelTemplate
 from tests.base.datamodules import TrialMNISTDataModule
+from tests.base import BoringModel, RandomDataset
+from torch.utils.data.sampler import SequentialSampler
+from torch.utils.data.dataloader import DataLoader
 
 
 def test_error_on_more_than_1_optimizer(tmpdir):
@@ -262,3 +266,29 @@ def test_suggestion_with_non_finite_values(tmpdir):
 
     assert before_lr == after_lr, \
         'Learning rate was altered because of non-finite loss values'
+
+
+@pytest.mark.timeout(1)
+def test_lr_finder_fails_fast_on_bad_config(tmpdir):
+    """ Test that misconfiguration of learning_rate or lr in model fails BEFORE lr optimization and not after it. """
+    import time
+    train = RandomDataset(32, 64)
+    context = 'spawn'
+    train = DataLoader(train, batch_size=32, num_workers=1, multiprocessing_context=context, shuffle=True)
+
+    class ExtendedBoringModel(BoringModel):
+        def train_dataloader(self):
+            return train
+    model = ExtendedBoringModel()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=2,
+        auto_lr_find=True,
+        deterministic=True,
+        overfit_batches=1,
+    )
+    train_data_loader = trainer.replace_sampler(train, SequentialSampler(train.dataset))
+
+    with pytest.raises(MisconfigurationException):
+        trainer.tune(model, train_dataloader=train_data_loader)