Addressed handling multiple datasets with the dataloader

Author: Hananeh Oliaei

src/electrai/configs/MP/config_resnet.yaml

@@ -1,19 +1,18 @@
 # Dataset / loader parameters
 data:
   _target_: electrai.dataloader.dataset.RhoRead
-  root: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/rho_gga/mp_filelist.txt
-  split_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/rho_gga/split_limit_22M.json
+  datasets:
+    - root: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/mp_filelist.txt
+      split_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/split.json
+      val_frac: 0.005
   precision: f32
   batch_size: 1
   train_workers: 8
   val_workers: 2
   pin_memory: false
-  val_frac: 0.005
   drop_last: false
   augmentation: false
   random_seed: 42
-  # downsample_label: 0
-  # downsample_data: 0
 
 # Model
 model:

src/electrai/configs/MP/config_resunet.yaml

@@ -1,19 +1,18 @@
 # Dataset / loader parameters
 data:
   _target_: electrai.dataloader.dataset.RhoRead
-  root: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/rho_gga/mp_filelist.txt
-  split_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/rho_gga/split_limit_22M.json
+  datasets:
+    - root: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/mp_filelist.txt
+      split_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/split.json
+      val_frac: 0.005
   precision: f32
   batch_size: 1
   train_workers: 8
   val_workers: 2
   pin_memory: false
-  val_frac: 0.005
   drop_last: false
   augmentation: false
   random_seed: 42
-  # downsample_label: 0
-  # downsample_data: 0
 
 # Model
 model:

src/electrai/dataloader/collate.py

@@ -7,5 +7,5 @@ def collate_fn(batch):
     try:
         return default_collate(batch)
     except RuntimeError:
-        x, y, index = zip(*batch, strict=True)
-        return list(x), list(y), list(index)
+        x, y, index, dataset_id = zip(*batch, strict=True)
+        return list(x), list(y), list(index), list(dataset_id)

src/electrai/dataloader/dataset.py

@@ -1,66 +1,141 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import Any
 
 import torch
 from lightning.pytorch import LightningDataModule
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import ConcatDataset, DataLoader, Dataset
 
 from electrai.dataloader import utils
 from electrai.dataloader.collate import collate_fn
 from electrai.dataloader.split import split_data
 
-if TYPE_CHECKING:
-    import os
-
 dtype_map = {"f32": torch.float32, "f16": torch.float16, "bf16": torch.bfloat16}
 
 
+@dataclass(frozen=True)
+class DatasetSpec:
+    root: str
+    split_file: str | None = None
+    val_frac: float | None = None
+    dataset_id: int | None = None
+
+
+class AddDatasetID(Dataset):
+    """Wrap dataset so every sample carries a constant Dataset_ID"""
+
+    def __init__(self, base: Dataset, functional_id: int, key: str = "Dataset_ID"):
+        self.base = base
+        self.functional_id = int(functional_id)
+        self.key = key
+
+    def __len__(self):
+        return len(self.base)
+
+    def __getitem__(self, idx):
+        out = self.base[idx]
+        out[self.key] = self.functional_id
+        return out
+
+
 class RhoRead(LightningDataModule):
+    """
+    Works based on these keys for one or more datasets:
+    `root`, `split_file`, `val_frac` (if `split_file` is null) and dataset_id
+    """
+
     def __init__(
         self,
-        root: str | bytes | os.PathLike,
-        precision: str,
+        datasets: list[dict[str, Any]] | list[DatasetSpec] | None = None,
+        default_val_frac: float = 0.005,
+        default_split_file: str | None = None,
+        precision: str = "f32",
         batch_size: int = 2,
         train_workers: int = 8,
         val_workers: int = 2,
         pin_memory: bool = False,
-        val_frac: float = 0.005,
         drop_last: bool = False,
-        split_file: str | bytes | os.PathLike | None = None,
         augmentation: bool = False,
         random_seed: int = 42,
     ):
         super().__init__()
         self.save_hyperparameters()
-        self.root = root
         self.batch_size = batch_size
         self.train_workers = train_workers
         self.val_workers = val_workers
         self.pin_memory = pin_memory
-        self.val_frac = val_frac
         self.drop_last = drop_last
-        self.split_file = split_file
         self.precision = precision
         self.augmentation = augmentation
         self.random_seed = random_seed
 
-    def setup(self, stage=None):
-        dataset = RhoData(
-            self.root, precision=self.precision, augmentation=self.augmentation
-        )
-        self.subsets = split_data(
-            dataset,
-            val_frac=self.val_frac,
-            split_file=self.split_file,
-            random_seed=self.random_seed,
-        )
+        if datasets is None or len(datasets) == 0:
+            raise ValueError("`datasets` must contain at least one dataset spec.")
+
+        specs: list[DatasetSpec] = [
+            d if isinstance(d, DatasetSpec) else DatasetSpec(**d) for d in datasets
+        ]
+        filled: list[DatasetSpec] = []
+        for i, s in enumerate(specs):
+            if s.root is None:
+                raise ValueError(
+                    f"`root` is required for dataset {i + 1}. Received: {s.root!r}"
+                )
+            dataset_id = s.dataset_id if s.dataset_id is not None else i
+            filled.append(
+                DatasetSpec(
+                    root=s.root,
+                    split_file=(
+                        s.split_file if s.split_file is not None else default_split_file
+                    ),
+                    val_frac=(
+                        s.val_frac if s.val_frac is not None else default_val_frac
+                    ),
+                    dataset_id=dataset_id,
+                )
+            )
+        self.specs = filled
+
+        self.train_set: Dataset | None = None
+        self.val_set: Dataset | None = None
+        self.test_set: Dataset | None = None
+
+    def setup(self, stage: str):
+        train_parts: list[Dataset] = []
+        val_parts: list[Dataset] = []
+        test_parts: list[Dataset] = []
+
+        for spec in self.specs:
+            ds = RhoData(
+                spec.root, precision=self.precision, augmentation=self.augmentation
+            )
+            splits = split_data(
+                ds,
+                val_frac=float(spec.val_frac),
+                split_file=spec.split_file,
+                random_seed=self.random_seed,
+            )
+
+            dataset_id = int(spec.dataset_id)
+
+            train_parts.append(AddDatasetID(splits["train"], dataset_id))
+            val_parts.append(AddDatasetID(splits["validation"], dataset_id))
+            if "test" in splits and splits["test"] is not None:
+                test_parts.append(AddDatasetID(splits["test"], dataset_id))
+
         if stage == "fit":
-            self.train_set = self.subsets["train"]
-            self.val_set = self.subsets["validation"]
+            self.train_set = (
+                train_parts[0] if len(train_parts) == 1 else ConcatDataset(train_parts)
+            )
+            self.val_set = (
+                val_parts[0] if len(val_parts) == 1 else ConcatDataset(val_parts)
+            )
         elif stage == "test":
-            self.test_set = self.subsets["test"]
+            self.test_set = (
+                test_parts[0] if len(test_parts) == 1 else ConcatDataset(test_parts)
+            )
 
     def train_dataloader(self):
         return DataLoader(