Adding flags to datamodules (Lightning-Universe#388)

* Adding flags to datamodules

* Finishing up changes

* Fixing syntax error

* More syntax errors

* More

* Adding drop_last flag to sklearn test

* Adding drop_last flag to sklearn test

* Updating doc for reflect drop_last=False

* Adding flags to datamodules

* Finishing up changes

* Fixing syntax error

* More syntax errors

* More

* Adding drop_last flag to sklearn test

* Adding drop_last flag to sklearn test

* Updating doc for reflect drop_last=False

* Cleaning up parameters and docstring

* Fixing syntax error

* Fixing documentation

* Hardcoding shuffle=False for val and test

pl_bolts/datamodules/binary_mnist_datamodule.py

@@ -48,8 +48,11 @@ def __init__(
             val_split: int = 5000,
             num_workers: int = 16,
             normalize: bool = False,
-            seed: int = 42,
             batch_size: int = 32,
+            seed: int = 42,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -60,6 +63,11 @@ def __init__(
             num_workers: how many workers to use for loading data
             normalize: If true applies image normalize
             batch_size: size of batch
+            seed: random seed to be used for train/val/test splits
+            shuffle: If true shuffles the data every epoch
+            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
+                        returning them
+            drop_last: If true drops the last incomplete batch
         """
         super().__init__(*args, **kwargs)
 
@@ -73,8 +81,11 @@ def __init__(
         self.val_split = val_split
         self.num_workers = num_workers
         self.normalize = normalize
-        self.seed = seed
         self.batch_size = batch_size
+        self.seed = seed
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
 
     @property
     def num_classes(self):
@@ -107,10 +118,10 @@ def train_dataloader(self):
         loader = DataLoader(
             dataset_train,
             batch_size=self.batch_size,
-            shuffle=True,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -131,8 +142,8 @@ def val_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -148,8 +159,8 @@ def test_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 

pl_bolts/datamodules/cifar10_datamodule.py

@@ -68,6 +68,9 @@ def __init__(
             num_workers: int = 16,
             batch_size: int = 32,
             seed: int = 42,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -77,6 +80,11 @@ def __init__(
             val_split: how many of the training images to use for the validation split
             num_workers: how many workers to use for loading data
             batch_size: number of examples per training/eval step
+            seed: random seed to be used for train/val/test splits
+            shuffle: If true shuffles the data every epoch
+            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
+                        returning them
+            drop_last: If true drops the last incomplete batch
         """
         super().__init__(*args, **kwargs)
 
@@ -91,6 +99,9 @@ def __init__(
         self.num_workers = num_workers
         self.batch_size = batch_size
         self.seed = seed
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
         self.data_dir = data_dir if data_dir is not None else os.getcwd()
         self.num_samples = 50000 - val_split
 
@@ -125,10 +136,10 @@ def train_dataloader(self):
         loader = DataLoader(
             dataset_train,
             batch_size=self.batch_size,
-            shuffle=True,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -150,8 +161,8 @@ def val_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            pin_memory=True,
-            drop_last=True
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last
         )
         return loader
 
@@ -167,8 +178,8 @@ def test_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 

pl_bolts/datamodules/cityscapes_datamodule.py

@@ -66,6 +66,9 @@ def __init__(
             num_workers: int = 16,
             batch_size: int = 32,
             seed: int = 42,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -77,6 +80,11 @@ def __init__(
             target_type: targets to use, either 'instance' or 'semantic'
             num_workers: how many workers to use for loading data
             batch_size: number of examples per training/eval step
+            seed: random seed to be used for train/val/test splits
+            shuffle: If true shuffles the data every epoch
+            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
+                        returning them
+            drop_last: If true drops the last incomplete batch
         """
         super().__init__(*args, **kwargs)
 
@@ -95,6 +103,9 @@ def __init__(
         self.num_workers = num_workers
         self.batch_size = batch_size
         self.seed = seed
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
         self.target_transforms = None
 
     @property
@@ -123,10 +134,10 @@ def train_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=True,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -150,8 +161,8 @@ def val_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            pin_memory=True,
-            drop_last=True
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last
         )
         return loader
 
@@ -174,8 +185,8 @@ def test_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 

pl_bolts/datamodules/fashion_mnist_datamodule.py

@@ -48,8 +48,11 @@ def __init__(
             data_dir: str,
             val_split: int = 5000,
             num_workers: int = 16,
-            seed: int = 42,
             batch_size: int = 32,
+            seed: int = 42,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -59,6 +62,11 @@ def __init__(
             val_split: how many of the training images to use for the validation split
             num_workers: how many workers to use for loading data
             batch_size: size of batch
+            seed: random seed to be used for train/val/test splits
+            shuffle: If true shuffles the data every epoch
+            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
+                        returning them
+            drop_last: If true drops the last incomplete batch
         """
         super().__init__(*args, **kwargs)
 
@@ -71,8 +79,11 @@ def __init__(
         self.data_dir = data_dir
         self.val_split = val_split
         self.num_workers = num_workers
-        self.seed = seed
         self.batch_size = batch_size
+        self.seed = seed
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
 
     @property
     def num_classes(self):
@@ -105,10 +116,10 @@ def train_dataloader(self):
         loader = DataLoader(
             dataset_train,
             batch_size=self.batch_size,
-            shuffle=True,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -130,8 +141,8 @@ def val_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -147,8 +158,8 @@ def test_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 

pl_bolts/datamodules/imagenet_datamodule.py

@@ -55,6 +55,9 @@ def __init__(
             image_size: int = 224,
             num_workers: int = 16,
             batch_size: int = 32,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -66,6 +69,10 @@ def __init__(
             image_size: final image size
             num_workers: how many data workers
             batch_size: batch_size
+            shuffle: If true shuffles the data every epoch
+            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
+                        returning them
+            drop_last: If true drops the last incomplete batch
         """
         super().__init__(*args, **kwargs)
 
@@ -81,6 +88,9 @@ def __init__(
         self.meta_dir = meta_dir
         self.num_imgs_per_val_class = num_imgs_per_val_class
         self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
         self.num_samples = 1281167 - self.num_imgs_per_val_class * self.num_classes
 
     @property
@@ -143,10 +153,10 @@ def train_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=True,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -170,7 +180,8 @@ def val_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -190,8 +201,8 @@ def test_dataloader(self):
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            drop_last=True,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader