PVNet 5.0.1__tar.gz → 5.0.3__tar.gz

{pvnet-5.0.1 → pvnet-5.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PVNet
-Version: 5.0.1
+Version: 5.0.3
 Summary: PVNet
 Author-email: Peter Dudfield <info@openclimatefix.org>
 Requires-Python: >=3.10

{pvnet-5.0.1 → pvnet-5.0.3}/PVNet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PVNet
-Version: 5.0.1
+Version: 5.0.3
 Summary: PVNet
 Author-email: Peter Dudfield <info@openclimatefix.org>
 Requires-Python: >=3.10

{pvnet-5.0.1 → pvnet-5.0.3}/pvnet/data/base_datamodule.py

@@ -46,6 +46,7 @@ class BasePresavedDataModule(LightningDataModule):
         num_workers: int = 0,
         prefetch_factor: int | None = None,
         persistent_workers: bool = False,
+        pin_memory: bool = False,
     ):
         """Base Datamodule for loading pre-saved samples
 
@@ -53,10 +54,12 @@ class BasePresavedDataModule(LightningDataModule):
             sample_dir: Path to the directory of pre-saved samples.
             batch_size: Batch size.
             num_workers: Number of workers to use in multiprocess batch loading.
-            prefetch_factor: Number of data will be prefetched at the end of each worker process.
+            prefetch_factor: Number of batches loaded in advance by each worker.
             persistent_workers: If True, the data loader will not shut down the worker processes 
                 after a dataset has been consumed once. This allows to maintain the workers Dataset 
                 instances alive.
+            pin_memory: If True, the data loader will copy Tensors into device/CUDA pinned memory 
+                before returning them.
         """
         super().__init__()
 
@@ -68,7 +71,7 @@ class BasePresavedDataModule(LightningDataModule):
             batch_sampler=None,
             num_workers=num_workers,
             collate_fn=collate_fn,
-            pin_memory=False,
+            pin_memory=pin_memory,
             drop_last=False,
             timeout=0,
             worker_init_fn=None,
@@ -100,8 +103,10 @@ class BaseStreamedDataModule(LightningDataModule):
         num_workers: int = 0,
         prefetch_factor: int | None = None,
         persistent_workers: bool = False,
+        pin_memory: bool = False,
         train_period: list[str | None] = [None, None],
         val_period: list[str | None] = [None, None],
+        seed: int | None = None,
 
     ):
         """Base Datamodule for streaming samples.
@@ -110,25 +115,29 @@ class BaseStreamedDataModule(LightningDataModule):
             configuration: Path to ocf-data-sampler configuration file.
             batch_size: Batch size.
             num_workers: Number of workers to use in multiprocess batch loading.
-            prefetch_factor: Number of data will be prefetched at the end of each worker process.
+            prefetch_factor: Number of batches loaded in advance by each worker.
             persistent_workers: If True, the data loader will not shut down the worker processes 
                 after a dataset has been consumed once. This allows to maintain the workers Dataset 
                 instances alive.
+            pin_memory: If True, the data loader will copy Tensors into device/CUDA pinned memory 
+                before returning them.
             train_period: Date range filter for train dataloader.
             val_period: Date range filter for val dataloader.
+            seed: Random seed used in shuffling datasets.
         """
         super().__init__()
 
         self.configuration = configuration
         self.train_period = train_period
         self.val_period = val_period
+        self.seed = seed
 
         self._common_dataloader_kwargs = dict(
             batch_size=batch_size,
             batch_sampler=None,
             num_workers=num_workers,
             collate_fn=collate_fn,
-            pin_memory=False,
+            pin_memory=pin_memory,
             drop_last=False,
             timeout=0,
             worker_init_fn=None,
@@ -142,11 +151,16 @@ class BaseStreamedDataModule(LightningDataModule):
         # This logic runs only once at the start of training, therefore the val dataset is only
         # shuffled once
         if stage == "fit":
+
             # Prepare the train dataset
             self.train_dataset = self._get_streamed_samples_dataset(*self.train_period)
 
-            # Prepare and pre-shuffle the val dataset
+            # Prepare and pre-shuffle the val dataset and set seed for reproducibility
             val_dataset = self._get_streamed_samples_dataset(*self.val_period)
+
+            if self.seed is not None:
+                torch.manual_seed(self.seed)
+
             shuffled_indices = torch.randperm(len(val_dataset))
             self.val_dataset = Subset(val_dataset, shuffled_indices)