Enhance DataLoader configurations in BlueSkyNonGeoDataModule by addin…

…g pin_memory, persistent_workers, and prefetch_factor parameters for improved performance during training, validation, and testing.
NRCan · Nov 28, 2024 · 60196fc · 60196fc
1 parent aadc64a
commit 60196fc
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/geo_deep_learning/datamodules/imagery_NonGeoDataModule.py b/geo_deep_learning/datamodules/imagery_NonGeoDataModule.py
@@ -73,17 +73,29 @@ def train_dataloader(self) -> DataLoader[Any]:
 
         return DataLoader(self.train_dataset, 
                           batch_size=self.batch_size, 
-                          num_workers=self.num_workers, shuffle=True)
+                          num_workers=self.num_workers,
+                          pin_memory=True,
+                          persistent_workers=True,
+                          prefetch_factor=2,
+                          shuffle=True)
 
     def val_dataloader(self):
         return DataLoader(self.val_dataset, 
                           batch_size=self.batch_size,
-                          num_workers=self.num_workers, shuffle=False)
+                          num_workers=self.num_workers,
+                          pin_memory=True,
+                          persistent_workers=True,
+                          prefetch_factor=2,
+                          shuffle=False)
 
     def test_dataloader(self):
         return DataLoader(self.test_dataset, 
                           batch_size=self.batch_size,
-                          num_workers=self.num_workers, shuffle=False)
+                          num_workers=self.num_workers,
+                          pin_memory=True,
+                          persistent_workers=True,
+                          prefetch_factor=2,
+                          shuffle=False)
 
     def _manage_bands(self, image: torch.Tensor) -> torch.Tensor:
         """