Adding flags to datamodules

pl_bolts/datamodules/binary_mnist_datamodule.py

@@ -53,6 +53,9 @@ def __init__(
             num_workers: int = 16,
             normalize: bool = False,
             seed: int = 42,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -62,6 +65,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
             normalize: If true applies image normalize
+            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)
 
@@ -76,6 +84,9 @@ def __init__(
         self.num_workers = num_workers
         self.normalize = normalize
         self.seed = seed
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
 
     @property
     def num_classes(self):
@@ -112,10 +123,10 @@ def train_dataloader(self, batch_size=32, transforms=None):
         loader = DataLoader(
             dataset_train,
             batch_size=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
 
@@ -138,10 +149,10 @@ def val_dataloader(self, batch_size=32, transforms=None):
         loader = DataLoader(
             dataset_val,
             batch_size=batch_size,
-            shuffle=False,
+            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
 
@@ -159,10 +170,10 @@ def test_dataloader(self, batch_size=32, transforms=None):
         loader = DataLoader(
             dataset,
             batch_size=batch_size,
-            shuffle=False,
+            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
 

pl_bolts/datamodules/cifar10_datamodule.py

@@ -71,6 +71,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,
     ):
@@ -80,6 +83,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)
 
@@ -94,6 +102,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 = 60000 - val_split
 
@@ -128,10 +139,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
 
@@ -151,10 +162,10 @@ def val_dataloader(self):
         loader = DataLoader(
             dataset_val,
             batch_size=self.batch_size,
-            shuffle=False,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            pin_memory=True,
-            drop_last=True
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last
         )
         return loader
 
@@ -168,10 +179,10 @@ def test_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=False,
+            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
 

pl_bolts/datamodules/cityscapes_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,
     ):
@@ -79,6 +82,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)
 
@@ -97,6 +105,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
@@ -125,10 +136,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,10 +161,10 @@ def val_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=False,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            pin_memory=True,
-            drop_last=True
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last
         )
         return loader
 
@@ -174,10 +185,10 @@ def test_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=False,
+            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
 

pl_bolts/datamodules/fashion_mnist_datamodule.py

@@ -51,6 +51,9 @@ def __init__(
             val_split: int = 5000,
             num_workers: int = 16,
             seed: int = 42,
+            shuffle: bool = False,
+            pin_memory: bool = False,
+            drop_last: bool = False,
             *args,
             **kwargs,
     ):
@@ -59,6 +62,11 @@ def __init__(
             data_dir: where to save/load the data
             val_split: how many of the training images to use for the validation split
             num_workers: how many workers to use for loading data
+            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)
 
@@ -72,6 +80,9 @@ def __init__(
         self.val_split = val_split
         self.num_workers = num_workers
         self.seed = seed
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
 
     @property
     def num_classes(self):
@@ -108,10 +119,10 @@ def train_dataloader(self, batch_size=32, transforms=None):
         loader = DataLoader(
             dataset_train,
             batch_size=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
 
@@ -135,10 +146,10 @@ def val_dataloader(self, batch_size=32, transforms=None):
         loader = DataLoader(
             dataset_val,
             batch_size=batch_size,
-            shuffle=False,
+            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
 
@@ -156,10 +167,10 @@ def test_dataloader(self, batch_size=32, transforms=None):
         loader = DataLoader(
             dataset,
             batch_size=batch_size,
-            shuffle=False,
+            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
 

pl_bolts/datamodules/imagenet_datamodule.py

@@ -64,6 +64,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,
     ):
@@ -74,6 +77,10 @@ def __init__(
             num_imgs_per_val_class: how many images per class for the validation set
             image_size: final image size
             num_workers: how many data workers
+            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
             batch_size: batch_size
         """
         super().__init__(*args, **kwargs)
@@ -90,6 +97,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
@@ -152,10 +162,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
 
@@ -177,9 +187,10 @@ def val_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=False,
+            shuffle=self.shuffle,
             num_workers=self.num_workers,
-            pin_memory=True
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory
         )
         return loader
 
@@ -197,10 +208,10 @@ def test_dataloader(self):
         loader = DataLoader(
             dataset,
             batch_size=self.batch_size,
-            shuffle=False,
+            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