Repository: Lightning-AI/pytorch-lightning

Issue:
DataModule with lr_find not supported
<!-- 
### 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 ``Trainer (trainer)``, a ``LightningModule (model)`` and a ``DataModule (data_module)``
2. Call ``trainer.lr_find(model, datamodule=data_module)``
3. Error: ``TypeError: lr_find() got an unexpected keyword argument 'datamodule'``

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
```
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader
class DataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
    def gen_set(self, start, end):
        vals = torch.tensor([float(i) for i in range(0, 800)])
        return torch.stack([torch.sin(vals), torch.cos(vals)], 1)
    def train_dataloader(self):
        return DataLoader(dataset=self.gen_set(0, 800))
    def val_dataloader(self):
        return DataLoader(dataset=self.gen_set(800, 900))
    def test_dataloader(self):
        return DataLoader(dataset=self.gen_set(900, 1000))

class LitModel(pl.LightningModule):
    def __init__(self, in_features):
        super().__init__()
        self.in_layer = nn.Linear(in_features, in_features//2)
        self.out_layer = nn.Linear(in_features//2, in_features)
        self.criterion = nn.MSELoss()
        self.lr = 0.001
    def forward(self, inp):
        inp = torch.tanh(self.in_layer(inp))
        return self.out_layer(inp)
    def training_step(self, batch, batch_idx):
        x = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, x.float())
        return pl.TrainResult(minimize=loss)
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=(self.lr or self.learning_rate))
    
data_module = DataModule()
trainer = pl.Trainer()
model = LitModel(in_features=2)
lr_finder = trainer.lr_find(model, datamodule=data_module)   #error
```

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
The ``lr_find()`` function should work as if it were passed a ``train_loader`` and ``test_loader`` when invoking it, just like ``fit()``

### Environment
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.5.1
	- pytorch-lightning: 0.9.0
	- tensorboard:       2.2.0
	- tqdm:              4.45.0
* System:
	- OS:                Linux
	- architecture:
		- 64bit
	- processor:         x86_64
	- python:            3.7.6


Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/pytorch_lightning/tuner/lr_finder.py b/pytorch_lightning/tuner/lr_finder.py
--- a/pytorch_lightning/tuner/lr_finder.py
+++ b/pytorch_lightning/tuner/lr_finder.py
@@ -11,21 +11,23 @@
 # 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.
+import importlib
 import os
+from typing import List, Optional, Sequence, Union
+
+import numpy as np
 import torch
-from typing import Optional, Sequence, List, Union
+from torch.optim.lr_scheduler import _LRScheduler
 from torch.utils.data import DataLoader
+
+from pytorch_lightning import _logger as log
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.core.datamodule import LightningDataModule
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.loggers.base import DummyLogger
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from torch.optim.lr_scheduler import _LRScheduler
-import importlib
-from pytorch_lightning import _logger as log
-import numpy as np
-from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.utilities.parsing import lightning_hasattr, lightning_setattr
 
-
 # check if ipywidgets is installed before importing tqdm.auto
 # to ensure it won't fail and a progress bar is displayed
 if importlib.util.find_spec('ipywidgets') is not None:
@@ -71,6 +73,7 @@ def lr_find(
         num_training: int = 100,
         mode: str = 'exponential',
         early_stop_threshold: float = 4.0,
+        datamodule: Optional[LightningDataModule] = None,
 ):
     r"""
     lr_find enables the user to do a range test of good initial learning rates,
@@ -81,7 +84,7 @@ def lr_find(
 
         train_dataloader: A PyTorch
             DataLoader with training samples. If the model has
-            a predefined train_dataloader method this will be skipped.
+            a predefined train_dataloader method, this will be skipped.
 
         min_lr: minimum learning rate to investigate
 
@@ -98,6 +101,12 @@ def lr_find(
             loss at any point is larger than early_stop_threshold*best_loss
             then the search is stopped. To disable, set to None.
 
+        datamodule: An optional `LightningDataModule` which holds the training
+            and validation dataloader(s). Note that the `train_dataloader` and
+            `val_dataloaders` parameters cannot be used at the same time as
+            this parameter, or a `MisconfigurationException` will be raised.
+
+
     Example::
 
         # Setup model and trainer
@@ -167,7 +176,8 @@ def lr_find(
     # Fit, lr & loss logged in callback
     trainer.fit(model,
                 train_dataloader=train_dataloader,
-                val_dataloaders=val_dataloaders)
+                val_dataloaders=val_dataloaders,
+                datamodule=datamodule)
 
     # Prompt if we stopped early
     if trainer.global_step != num_training:
diff --git a/pytorch_lightning/tuner/tuning.py b/pytorch_lightning/tuner/tuning.py
--- a/pytorch_lightning/tuner/tuning.py
+++ b/pytorch_lightning/tuner/tuning.py
@@ -15,6 +15,7 @@
 from pytorch_lightning.tuner.auto_gpu_select import pick_multiple_gpus
 from pytorch_lightning.tuner.lr_finder import _run_lr_finder_internally, lr_find
 from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.core.datamodule import LightningDataModule
 from typing import Optional, List, Union
 from torch.utils.data import DataLoader
 
@@ -50,6 +51,7 @@ def lr_find(
             num_training: int = 100,
             mode: str = 'exponential',
             early_stop_threshold: float = 4.0,
+            datamodule: Optional[LightningDataModule] = None
     ):
         return lr_find(
             self.trainer,
@@ -60,7 +62,8 @@ def lr_find(
             max_lr,
             num_training,
             mode,
-            early_stop_threshold
+            early_stop_threshold,
+            datamodule,
         )
 
     def internal_find_lr(self, trainer, model: LightningModule):


Test patch:
diff --git a/tests/trainer/test_lr_finder.py b/tests/trainer/test_lr_finder.py
--- a/tests/trainer/test_lr_finder.py
+++ b/tests/trainer/test_lr_finder.py
@@ -5,6 +5,7 @@
 from pytorch_lightning import Trainer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
+from tests.base.datamodules import TrialMNISTDataModule
 
 
 def test_error_on_more_than_1_optimizer(tmpdir):
@@ -152,6 +153,30 @@ def test_call_to_trainer_method(tmpdir):
         'Learning rate was not altered after running learning rate finder'
 
 
+def test_datamodule_parameter(tmpdir):
+    """ Test that the datamodule parameter works """
+
+    # trial datamodule
+    dm = TrialMNISTDataModule(tmpdir)
+
+    hparams = EvalModelTemplate.get_default_hparams()
+    model = EvalModelTemplate(**hparams)
+
+    before_lr = hparams.get('learning_rate')
+    # logger file to get meta
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=2,
+    )
+
+    lrfinder = trainer.tuner.lr_find(model, datamodule=dm)
+    after_lr = lrfinder.suggestion()
+    model.learning_rate = after_lr
+
+    assert before_lr != after_lr, \
+        'Learning rate was not altered after running learning rate finder'
+
+
 def test_accumulation_and_early_stopping(tmpdir):
     """ Test that early stopping of learning rate finder works, and that
         accumulation also works for this feature """