new way of passing dataloaders

pytorch_lightning/core/decorators.py

@@ -9,9 +9,8 @@ def data_loader(fn):
     :return:
     """
 
-    wraps(fn)
     attr_name = '_lazy_' + fn.__name__
-
+    @wraps(fn)
     def _get_data_loader(self):
         try:
             value = getattr(self, attr_name)

pytorch_lightning/trainer/trainer.py

@@ -19,7 +19,7 @@
     parse_gpu_ids,
     determine_root_gpu_device
 )
-
+from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.trainer.evaluation_loop import TrainerEvaluationLoopMixin
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.trainer.model_hooks import TrainerModelHooksMixin
@@ -713,17 +713,38 @@ def tng_tqdm_dic(self):
     # -----------------------------
     # MODEL TRAINING
     # -----------------------------
-    def fit(self, model):
+    def fit(self, model, train_dataloader=None, val_dataloader=None, test_dataloader=None):
         r"""
         Runs the full optimization routine.
 
-        Example::
+        Args:
+            model (LightningModule): Model to fit.
+            Example::
 
-            trainer = Trainer()
-            model = LightningModule()
+                trainer = Trainer()
+                model = LightningModule()
+                trainer.fit(model)
+
+            train_dataloader (:class:`.torch.utils.data.DataLoader`): A Pytorch
+                DataLoader with training samples. If the model has
+                a predefined train_dataloader method this will be skipped.
+
+            val_dataloader (:class:`.torch.utils.data.DataLoader`): Either a single
+                Pytorch Dataloader or a list of them, specifying validation samples. If the model has
+                a predefined val_dataloader method this will be skipped
+
+            test_dataloader (:class:`.torch.utils.data.DataLoader`): Either a single
+                Pytorch Dataloader or a list of them, specifying validation samples. If the model has
+                a predefined val_dataloader method this will be skipped
 
-            trainer.fit()
         """
+
+        # Update the dataloader attributes of the model with the ones supplied here,
+        # if they are not already defined in model
+        _set_dataloader(model, train_dataloader, 'train_dataloader')
+        _set_dataloader(model, val_dataloader, 'val_dataloader')
+        _set_dataloader(model, test_dataloader, 'test_dataloader')
+
         # when using multi-node or DDP within a node start each module in a separate process
         if self.use_ddp2:
             task = int(os.environ['SLURM_LOCALID'])
@@ -911,3 +932,34 @@ def test(self, model=None):
             self.fit(model)
         else:
             self.run_evaluation(test=True)
+
+
+def _set_dataloader(model, dataloader, attribute):
+    r''' Check dataloaders passed to .fit() method if they are pytorch DataLoader
+         objects and whether or not we should overright the corresponding dataloader
+         in the model '''
+
+    # Check if attribute comes directly from base class or
+    # derived in user subclass
+    if LightningModule.__qualname__ in getattr(model, attribute).__qualname__:
+        # Val and test should be list of dataloaders
+        dataloader = dataloader if attribute == 'train_dataloader' or \
+            (attribute != 'train_dataloader' and isinstance(dataloader, list)) else [dataloader]
+
+        # Check if input is correct
+        if isinstance(dataloader, torch.utils.data.DataLoader) or \
+           isinstance(dataloader, list) and all(isinstance(d, torch.utils.data.DataLoader) for d in dataloader):
+
+            # Overwrite abstract methods
+            dl = lambda: dataloader
+            dl.__name__ = attribute
+            setattr(model, attribute, dl)
+
+        elif dataloader:
+            raise ValueError(f'`{attribute}` needs to be an instance of'
+                             '`torch.utils.data.DataLoader` or a list of, instead got'
+                             f'`{dataloader}`')
+
+    else:
+        logging.info(f'Model has predefined `{attribute}`,'
+                     f'will skip `{attribute}` passed to fit method')