Repository: Lightning-AI/pytorch-lightning

Issue:
JIT support
**Initial issue description below**

JIT support requires several changes in the `Trainer` and `LightningModule`:
- [x] No use of python properties like the current dataloader implementation
    - Possible solution: Use getters like implemented in #275 
    - Other possibility: Handle dataloading completely in trainer. The user is able to modify the dataloaders e.g. every epoch using hooks/callbacks
- [ ] The trainer cannot set PLModule's class members afterwards like `self.model.trainer = self`
    - This is because after converting `my_pl_module = torch.jit.script(my_pl_module)`, the module has the class `ScriptModule`. Adding members only adds the members to the `ScriptModule` not to the underlying `LightningModule`.
    - Solution could be: Implement setter in `LightningModule`. These methods will be transfered to the `ScriptModule`
- [ ] Saving and restoring might need some changes, too. One could conditionally check the class of the provided module in the trainer for use of `torch.jit.save` in trainer_io
- [ ]  JIT is currently not compatible with distributed training (see pytorch issue [#15421](https://github.com/pytorch/pytorch/issues/15421))



**Is your feature request related to a problem? Please describe.**
The current implementation of `pl.LightningModule` does not support pytorch's JIT. This is due to the use of python properties for the dataloaders, which is currently not supported in JIT ([see here](https://github.com/pytorch/pytorch/issues/23958)).

example trace:
```py
$ python train.py
VISIBLE GPUS:
Traceback (most recent call last):
  File "train.py", line 69, in <module>
    train(config, data_dir)
  File "train.py", line 37, in train
    trainer.fit(pl_module)
  File "/home/schroeter/.virtualenvs/pytorch-1.2/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 694, in fit
    self.__run_pretrain_routine(model)
  File "/home/schroeter/.virtualenvs/pytorch-1.2/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 879, in __run_pretrain_routine
    self.get_dataloaders(ref_model)
  File "/home/schroeter/.virtualenvs/pytorch-1.2/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 574, in get_dataloaders
    self.train_dataloader = model.train_dataloader
  File "/home/schroeter/.virtualenvs/pytorch-1.2/lib/python3.7/site-packages/torch/jit/__init__.py", line 1563, in __getattr__
    return super(ScriptModule, self).__getattr__(attr)
  File "/home/schroeter/.virtualenvs/pytorch-1.2/lib/python3.7/site-packages/torch/nn/modules/module.py", line 589, in __getattr__
    type(self).__name__, name))
AttributeError: 'ScriptModule' object has no attribute 'train_dataloader'
```
where pl_module has:
```py
class PlModule(pl.LightningModule):
# ....
    @pl.data_loader
    def train_dataloader(self):
        return self.__dataloader("train")
```

**Describe the solution you'd like**
```py
torch.jit.script(pl_module)
```

**Describe alternatives you've considered**
A workaround might be defining a separate Module handling all nn.Module stuff and transforming only this part into a jit script module.

A solution might be to define explicit getters like `get_train_dataloader(self)` instead of using properties.



Discussion:
Good point. Hadn't tested lightning with JIT. I think the choice of property vs getter made sense for the first design of the framework. However, recent refactors have made those differences irrelevant. 

I don't see any reason why we shouldn't move to a getter instead of the property. Do you want to submit a PR for this?
Yes I can do that

Resolving patch:
diff --git a/pytorch_lightning/root_module/decorators.py b/pytorch_lightning/root_module/decorators.py
--- a/pytorch_lightning/root_module/decorators.py
+++ b/pytorch_lightning/root_module/decorators.py
@@ -10,13 +10,18 @@ def data_loader(fn):
 
     attr_name = '_lazy_' + fn.__name__
 
-    @property
-    def _data_loader(self):
+    def _get_data_loader(self):
         try:
             value = getattr(self, attr_name)
         except AttributeError:
             try:
                 value = fn(self)  # Lazy evaluation, done only once.
+                if (
+                    value is not None and
+                    not isinstance(value, list) and
+                    fn.__name__ in['test_dataloader', 'val_dataloader']
+                ):
+                    value = [value]
             except AttributeError as e:
                 # Guard against AttributeError suppression. (Issue #142)
                 traceback.print_exc()
@@ -25,4 +30,4 @@ def _data_loader(self):
             setattr(self, attr_name, value)  # Memoize evaluation.
         return value
 
-    return _data_loader
+    return _get_data_loader
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -24,6 +24,7 @@
 import pdb
 from pytorch_lightning.trainer import ignored_warnings
 
+
 try:
     from apex import amp
     APEX_AVAILABLE = True
@@ -141,9 +142,9 @@ def __init__(self,
         self.nb_val_batches = 0
         self.nb_training_batches = 0
         self.nb_test_batches = 0
-        self.train_dataloader = None
-        self.test_dataloader = None
-        self.val_dataloader = None
+        self.get_train_dataloader = None
+        self.get_test_dataloaders = None
+        self.get_val_dataloaders = None
 
         # training state
         self.model = None
@@ -441,19 +442,21 @@ def training_tqdm_dict(self):
     def __layout_bookeeping(self):
 
         # determine number of training batches
-        self.nb_training_batches = len(self.train_dataloader)
+        self.nb_training_batches = len(self.get_train_dataloader())
         self.nb_training_batches = int(self.nb_training_batches * self.train_percent_check)
 
         # determine number of validation batches
         # val datasets could be none, 1 or 2+
-        if self.val_dataloader is not None:
-            self.nb_val_batches = sum(len(dataloader) for dataloader in self.val_dataloader)
+        if self.get_val_dataloaders() is not None:
+            self.nb_val_batches = sum(len(dataloader) for dataloader in self.get_val_dataloaders())
             self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
             self.nb_val_batches = max(1, self.nb_val_batches)
 
         # determine number of test batches
-        if self.test_dataloader is not None:
-            self.nb_test_batches = sum(len(dataloader) for dataloader in self.test_dataloader)
+        if self.get_test_dataloaders() is not None:
+            self.nb_test_batches = sum(
+                len(dataloader) for dataloader in self.get_test_dataloaders()
+            )
             self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
             self.nb_test_batches = max(1, self.nb_test_batches)
 
@@ -472,10 +475,10 @@ def __evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test=Fal
         # make dataloader_idx arg in validation_step optional
         args = [batch, batch_idx]
 
-        if test and len(self.test_dataloader) > 1:
+        if test and len(self.get_test_dataloaders()) > 1:
             args.append(dataloader_idx)
 
-        elif not test and len(self.val_dataloader) > 1:
+        elif not test and len(self.get_val_dataloaders()) > 1:
             args.append(dataloader_idx)
 
         # handle DP, DDP forward
@@ -520,9 +523,9 @@ def evaluate(self, model, dataloaders, max_batches, test=False):
         outputs = []
 
         # run training
-        for dataloader_idx, dl in enumerate(dataloaders):
+        for dataloader_idx, dataloader in enumerate(dataloaders):
             dl_outputs = []
-            for batch_idx, batch in enumerate(dl):
+            for batch_idx, batch in enumerate(dataloader):
 
                 if batch is None:  # pragma: no cover
                     continue
@@ -570,21 +573,11 @@ def get_dataloaders(self, model):
         :param model:
         :return:
         """
+        self.get_train_dataloader = model.train_dataloader
+        self.get_test_dataloaders = model.test_dataloader
+        self.get_val_dataloaders = model.val_dataloader
 
-        self.train_dataloader = model.train_dataloader
-        self.test_dataloader = model.test_dataloader
-        self.val_dataloader = model.val_dataloader
-
-        # handle returning an actual dataloader instead of a list of loaders
-        have_test_loaders = self.test_dataloader is not None
-        if have_test_loaders and not isinstance(self.test_dataloader, list):
-            self.test_dataloader = [self.test_dataloader]
-
-        have_val_loaders = self.val_dataloader is not None
-        if have_val_loaders and not isinstance(self.val_dataloader, list):
-            self.val_dataloader = [self.val_dataloader]
-
-        if self.use_ddp and not isinstance(self.train_dataloader.sampler, DistributedSampler):
+        if self.use_ddp and not isinstance(self.get_train_dataloader().sampler, DistributedSampler):
             msg = """
             You're using multiple gpus and multiple nodes without using a DistributedSampler
             to assign a subset of your data to each process. To silence this warning, pass a
@@ -603,8 +596,8 @@ def get_dataloaders(self, model):
             """
             warnings.warn(msg)
 
-        if self.use_ddp and self.val_dataloader is not None:
-            for dataloader in self.val_dataloader:
+        if self.use_ddp and self.get_val_dataloaders is not None:
+            for dataloader in self.get_val_dataloaders():
                 if not isinstance(dataloader.sampler, DistributedSampler):
                     msg = """
                     Your val_dataloader(s) don't use DistributedSampler.
@@ -626,8 +619,8 @@ def get_dataloaders(self, model):
                     warnings.warn(msg)
                     break
 
-        if self.use_ddp and self.test_dataloader is not None:
-            for dataloader in self.test_dataloader:
+        if self.use_ddp and self.get_test_dataloaders is not None:
+            for dataloader in self.get_test_dataloaders():
                 if not isinstance(dataloader.sampler, DistributedSampler):
                     msg = """
                     Your test_dataloader(s) don't use DistributedSampler.
@@ -912,12 +905,12 @@ def __run_pretrain_routine(self, model):
         # run tiny validation (if validation defined)
         # to make sure program won't crash during val
         ref_model.on_sanity_check_start()
-        if self.val_dataloader is not None and self.nb_sanity_val_steps > 0:
+        if self.get_val_dataloaders() is not None and self.nb_sanity_val_steps > 0:
             # reset progress_bar limit for sanity check
             if self.show_progress_bar:
                 self.progress_bar.reset(self.nb_sanity_val_steps)
 
-            self.evaluate(model, self.val_dataloader, self.nb_sanity_val_steps, self.testing)
+            self.evaluate(model, self.get_val_dataloaders(), self.nb_sanity_val_steps, self.testing)
 
         # ---------------------------
         # CORE TRAINING LOOP
@@ -928,8 +921,8 @@ def __train(self):
         # run all epochs
         for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
             # set seed for distributed sampler (enables shuffling for each epoch)
-            if self.use_ddp and hasattr(self.train_dataloader.sampler, 'set_epoch'):
-                self.train_dataloader.sampler.set_epoch(epoch_nb)
+            if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):
+                self.get_train_dataloader().sampler.set_epoch(epoch_nb)
 
             # get model
             model = self.__get_model()
@@ -974,7 +967,7 @@ def run_training_epoch(self):
             model.on_epoch_start()
 
         # run epoch
-        for batch_nb, batch in enumerate(self.train_dataloader):
+        for batch_nb, batch in enumerate(self.get_train_dataloader()):
             self.batch_nb = batch_nb
             self.global_step += 1
 
@@ -1272,12 +1265,12 @@ def __run_evaluation(self, test=False):
             model.on_pre_performance_check()
 
             # select dataloaders
-            dataloaders = self.val_dataloader
+            dataloaders = self.get_val_dataloaders()
             max_batches = self.nb_val_batches
 
             # calculate max batches to use
             if test:
-                dataloaders = self.test_dataloader
+                dataloaders = self.get_test_dataloaders()
                 max_batches = self.nb_test_batches
 
             # cap max batches to 1 when using fast_dev_run


Test patch:
diff --git a/tests/test_models.py b/tests/test_models.py
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -105,7 +105,8 @@ def assert_good_acc():
         dp_model = new_trainer.model
         dp_model.eval()
 
-        _ = [run_prediction(dataloader, dp_model, dp=True) for dataloader in trainer.val_dataloader]
+        for dataloader in trainer.get_train_dataloader():
+            run_prediction(dataloader, dp_model, dp=True)
 
     # new model
     model = LightningTestModel(hparams)
@@ -163,7 +164,7 @@ def test_running_test_pretrained_model_ddp():
     new_trainer = Trainer(**trainer_options)
     new_trainer.test(pretrained_model)
 
-    run_prediction(model.test_dataloader, pretrained_model)
+    [run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
 
     # test we have good test accuracy
     clear_save_dir()
@@ -761,7 +762,8 @@ def assert_good_acc():
         # if model and state loaded correctly, predictions will be good even though we
         # haven't trained with the new loaded model
         trainer.model.eval()
-        _ = [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
+        for dataloader in trainer.get_val_dataloaders():
+            run_prediction(dataloader, trainer.model)
 
     model.on_sanity_check_start = assert_good_acc
 
@@ -829,8 +831,9 @@ def test_cpu_slurm_save_load():
 
     # predict with trained model before saving
     # make a prediction
-    for batch in model.test_dataloader:
-        break
+    for dataloader in model.test_dataloader():
+        for batch in dataloader:
+            break
 
     x, y = batch
     x = x.view(x.size(0), -1)
@@ -945,8 +948,9 @@ def test_model_saving_loading():
     assert result == 1, 'amp + ddp model failed to complete'
 
     # make a prediction
-    for batch in model.test_dataloader:
-        break
+    for dataloader in model.test_dataloader():
+        for batch in dataloader:
+            break
 
     x, y = batch
     x = x.view(x.size(0), -1)
@@ -1037,7 +1041,7 @@ def test_amp_gpu_ddp_slurm_managed():
     pretrained_model = load_model(logger.experiment, save_dir, True)
 
     # test model preds
-    run_prediction(model.test_dataloader, pretrained_model)
+    [run_prediction(dataloader, pretrained_model) for dataloader in trainer.get_test_dataloaders()]
 
     if trainer.use_ddp:
         # on hpc this would work fine... but need to hack it for the purpose of the test
@@ -1264,10 +1268,11 @@ class CurrentTestModel(
     assert result == 1
 
     # verify there are 2 val loaders
-    assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly'
+    assert len(trainer.get_val_dataloaders()) == 2, \
+        'Multiple val_dataloaders not initiated properly'
 
     # make sure predictions are good for each val set
-    [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
+    [run_prediction(dataloader, trainer.model) for dataloader in trainer.get_val_dataloaders()]
 
 
 def test_multiple_test_dataloader():
@@ -1295,10 +1300,11 @@ class CurrentTestModel(
     result = trainer.fit(model)
 
     # verify there are 2 val loaders
-    assert len(trainer.test_dataloader) == 2, 'Multiple test_dataloaders not initiated properly'
+    assert len(trainer.get_test_dataloaders()) == 2, \
+        'Multiple test_dataloaders not initiated properly'
 
     # make sure predictions are good for each test set
-    [run_prediction(dataloader, trainer.model) for dataloader in trainer.test_dataloader]
+    [run_prediction(dataloader, trainer.model) for dataloader in trainer.get_test_dataloaders()]
 
     # run the test method
     trainer.test()
@@ -1333,7 +1339,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     pretrained_model = load_model(logger.experiment, save_dir, on_gpu)
 
     # test new model accuracy
-    run_prediction(model.test_dataloader, pretrained_model)
+    [run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
 
     if trainer.use_ddp:
         # on hpc this would work fine... but need to hack it for the purpose of the test