Fix for data_loader in pl version 0.7

Author: AntixK

experiment.py

@@ -3,14 +3,14 @@
 from torch import optim
 from models import BaseVAE
 from models.types_ import *
+from utils import data_loader
 import pytorch_lightning as pl
 from torchvision import transforms
 import torchvision.utils as vutils
 from torchvision.datasets import CelebA
 from torch.utils.data import DataLoader
 
 
-
 class VAEXperiment(pl.LightningModule):
 
     def __init__(self,
@@ -50,7 +50,7 @@ def validation_step(self, batch, batch_idx, optimizer_idx = 0):
 
         results = self.forward(real_img, labels = labels)
         val_loss = self.model.loss_function(*results,
-                                            M_N = self.params['batch_size']/ self.num_train_imgs,
+                                            M_N = self.params['batch_size']/ self.num_val_imgs,
                                             optimizer_idx = optimizer_idx,
                                             batch_idx = batch_idx)
 
@@ -132,7 +132,7 @@ def configure_optimizers(self):
         except:
             return optims
 
-    @pl.data_loader
+    @data_loader
     def train_dataloader(self):
         transform = self.data_transforms()
 
@@ -150,7 +150,7 @@ def train_dataloader(self):
                           shuffle = True,
                           drop_last=True)
 
-    @pl.data_loader
+    @data_loader
     def val_dataloader(self):
         transform = self.data_transforms()
 
@@ -162,8 +162,10 @@ def val_dataloader(self):
                                                  batch_size= 144,
                                                  shuffle = True,
                                                  drop_last=True)
+            self.num_val_imgs = len(self.sample_dataloader)
         else:
             raise ValueError('Undefined dataset type')
+
         return self.sample_dataloader
 
     def data_transforms(self):

utils.py

@@ -0,0 +1,22 @@
+import pytorch_lightning as pl
+
+
+## Utils to handle newer PyTorch Lightning changes from version 0.6
+## ==================================================================================================== ##
+
+
+def data_loader(fn):
+    """
+    Decorator to handle the deprecation of data_loader from 0.7
+    :param fn: User defined data loader function
+    :return: A wrapper for the data_loader function
+    """
+
+    def func_wrapper(self):
+        try: # Works for version 0.6.0
+            return pl.data_loader(fn)(self)
+
+        except: # Works for version > 0.6.0
+            return fn(self)
+
+    return func_wrapper