Adding flags to datamodules

pl_bolts/datamodules/binary_mnist_datamodule.py

@@ -52,8 +52,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,
     ):
@@ -64,6 +67,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)
 
@@ -77,8 +85,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):
@@ -111,10 +122,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
 
@@ -133,10 +144,10 @@ def val_dataloader(self):
         loader = DataLoader(
             dataset_val,
             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
 
@@ -150,10 +161,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/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 = 50000 - 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

@@ -50,8 +50,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,
     ):
@@ -61,6 +64,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)
 
@@ -73,8 +81,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):
@@ -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
 
@@ -130,10 +141,10 @@ def val_dataloader(self):
         loader = DataLoader(
             dataset_val,
             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
 
@@ -147,10 +158,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/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,
     ):
@@ -75,6 +78,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)
 
@@ -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