Use batch transformation functions from ocf_datapipes

Author: markus-kreft

pvnet/data/datamodule.py

@@ -3,12 +3,11 @@
 import resource
 
 import torch
-from ocf_datapipes.batch import stack_np_examples_into_batch
+from ocf_datapipes.batch import batch_to_tensor, stack_np_examples_into_batch
 from ocf_datapipes.training.pvnet import pvnet_datapipe
 from torch.utils.data.datapipes.iter import FileLister
 
 from pvnet.data.base import BaseDataModule
-from pvnet.data.utils import batch_to_tensor
 
 rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
 resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))

pvnet/data/pv_site_datamodule.py

@@ -1,11 +1,10 @@
 """ Data module for pytorch lightning """
 import glob
 
-from ocf_datapipes.batch import BatchKey, stack_np_examples_into_batch
+from ocf_datapipes.batch import BatchKey, batch_to_tensor, stack_np_examples_into_batch
 from ocf_datapipes.training.pvnet_site import pvnet_site_netcdf_datapipe
 
 from pvnet.data.base import BaseDataModule
-from pvnet.data.utils import batch_to_tensor
 
 
 class PVSiteDataModule(BaseDataModule):

pvnet/data/utils.py

@@ -1,36 +1,8 @@
 """Utils common between Wind and PV datamodules"""
-import numpy as np
-import torch
 from ocf_datapipes.batch import BatchKey, unstack_np_batch_into_examples
 from torch.utils.data import IterDataPipe, functional_datapipe
 
 
-def copy_batch_to_device(batch, device):
-    """Moves a dict-batch of tensors to new device."""
-    batch_copy = {}
-
-    for k, v in batch.items():
-        if isinstance(v, dict):
-            # Recursion to reach the nested NWP
-            batch_copy[k] = copy_batch_to_device(v, device)
-        elif isinstance(v, torch.Tensor):
-            batch_copy[k] = v.to(device)
-        else:
-            batch_copy[k] = v
-    return batch_copy
-
-
-def batch_to_tensor(batch):
-    """Moves numpy batch to a tensor"""
-    for k, v in batch.items():
-        if isinstance(v, dict):
-            # Recursion to reach the nested NWP
-            batch[k] = batch_to_tensor(v)
-        elif isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.number):
-            batch[k] = torch.as_tensor(v)
-    return batch
-
-
 @functional_datapipe("split_batches")
 class BatchSplitter(IterDataPipe):
     """Pipeline step to split batches of data and yield single examples"""

pvnet/data/wind_datamodule.py

@@ -1,11 +1,10 @@
 """ Data module for pytorch lightning """
 import glob
 
-from ocf_datapipes.batch import BatchKey, stack_np_examples_into_batch
+from ocf_datapipes.batch import BatchKey, batch_to_tensor, stack_np_examples_into_batch
 from ocf_datapipes.training.windnet import windnet_netcdf_datapipe
 
 from pvnet.data.base import BaseDataModule
-from pvnet.data.utils import batch_to_tensor
 
 
 class WindDataModule(BaseDataModule):

requirements.txt

@@ -1,4 +1,4 @@
-ocf_datapipes>=3.1.5
+ocf_datapipes>=3.3.6
 ocf_ml_metrics>=0.0.11
 numpy
 pandas

scripts/save_batches.py

@@ -33,7 +33,7 @@
 
 import hydra
 import torch
-from ocf_datapipes.batch import stack_np_examples_into_batch
+from ocf_datapipes.batch import batch_to_tensor, stack_np_examples_into_batch
 from ocf_datapipes.training.pvnet import pvnet_datapipe
 from ocf_datapipes.training.pvnet_site import pvnet_site_datapipe
 from ocf_datapipes.training.windnet import windnet_datapipe
@@ -43,7 +43,6 @@
 from torch.utils.data.datapipes.iter import IterableWrapper
 from tqdm import tqdm
 
-from pvnet.data.utils import batch_to_tensor
 from pvnet.utils import print_config
 
 warnings.filterwarnings("ignore", category=sa_exc.SAWarning)

scripts/save_concurrent_batches.py

@@ -23,7 +23,7 @@
 import hydra
 import numpy as np
 import torch
-from ocf_datapipes.batch import BatchKey, stack_np_examples_into_batch
+from ocf_datapipes.batch import BatchKey, batch_to_tensor, stack_np_examples_into_batch
 from ocf_datapipes.training.common import (
     open_and_return_datapipes,
 )
@@ -34,7 +34,6 @@
 from torch.utils.data.datapipes.iter import IterableWrapper
 from tqdm import tqdm
 
-from pvnet.data.utils import batch_to_tensor
 from pvnet.utils import GSPLocationLookup
 
 warnings.filterwarnings("ignore", category=sa_exc.SAWarning)