From c9dd6e36a4d9e675dff7dcbff965d87b5df7539e Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Fri, 13 Mar 2026 19:13:10 -0400 Subject: [PATCH 1/6] Addressed handling multiple datasets with the dataloader --- src/electrai/configs/MP/config_resnet.yaml | 9 +- src/electrai/configs/MP/config_resunet.yaml | 9 +- src/electrai/dataloader/collate.py | 4 +- src/electrai/dataloader/dataset.py | 125 ++++++++++++++++---- 4 files changed, 110 insertions(+), 37 deletions(-) diff --git a/src/electrai/configs/MP/config_resnet.yaml b/src/electrai/configs/MP/config_resnet.yaml index 482c48a9..406f3085 100644 --- a/src/electrai/configs/MP/config_resnet.yaml +++ b/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: diff --git a/src/electrai/configs/MP/config_resunet.yaml b/src/electrai/configs/MP/config_resunet.yaml index b270db22..30b425dc 100644 --- a/src/electrai/configs/MP/config_resunet.yaml +++ b/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: diff --git a/src/electrai/dataloader/collate.py b/src/electrai/dataloader/collate.py index 144924fa..e74c360b 100644 --- a/src/electrai/dataloader/collate.py +++ b/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) diff --git a/src/electrai/dataloader/dataset.py b/src/electrai/dataloader/dataset.py index 0a0a1c10..a473c2da 100644 --- a/src/electrai/dataloader/dataset.py +++ b/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( From 071e1d8c2e7ffb0f5a08b059ba612441c8fe6f70 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Tue, 19 May 2026 17:06:03 -0400 Subject: [PATCH 2/6] Address PR review comments - Fix collate_fn fallback to work with dict batches and return consistent dict - Gate train/val part building in setup() by stage - Remove dead `if s.root is None` guard - Add comment in lightning.py about Dataset_ID for future multi-head use - Add val_frac comment in configs clarifying it is ignored when split_file is set - Add tests for AddDatasetID, collate_fn fallback, DatasetSpec, RhoRead init and setup Co-Authored-By: Claude Sonnet 4.6 --- src/electrai/configs/MP/config_resnet.yaml | 2 +- src/electrai/configs/MP/config_resunet.yaml | 2 +- src/electrai/dataloader/collate.py | 3 +- src/electrai/dataloader/dataset.py | 14 +- src/electrai/lightning.py | 1 + tests/electrai/dataloader/test_dataset.py | 298 +++++++++++++++++++- 6 files changed, 307 insertions(+), 13 deletions(-) diff --git a/src/electrai/configs/MP/config_resnet.yaml b/src/electrai/configs/MP/config_resnet.yaml index 406f3085..244444fa 100644 --- a/src/electrai/configs/MP/config_resnet.yaml +++ b/src/electrai/configs/MP/config_resnet.yaml @@ -4,7 +4,7 @@ data: 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 + val_frac: 0.005 # ignored when split_file is provided precision: f32 batch_size: 1 train_workers: 8 diff --git a/src/electrai/configs/MP/config_resunet.yaml b/src/electrai/configs/MP/config_resunet.yaml index 30b425dc..3825d8c4 100644 --- a/src/electrai/configs/MP/config_resunet.yaml +++ b/src/electrai/configs/MP/config_resunet.yaml @@ -4,7 +4,7 @@ data: 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 + val_frac: 0.005 # ignored when split_file is provided precision: f32 batch_size: 1 train_workers: 8 diff --git a/src/electrai/dataloader/collate.py b/src/electrai/dataloader/collate.py index e74c360b..df5f31fb 100644 --- a/src/electrai/dataloader/collate.py +++ b/src/electrai/dataloader/collate.py @@ -7,5 +7,4 @@ def collate_fn(batch): try: return default_collate(batch) except RuntimeError: - x, y, index, dataset_id = zip(*batch, strict=True) - return list(x), list(y), list(index), list(dataset_id) + return {k: [d[k] for d in batch] for k in batch[0]} diff --git a/src/electrai/dataloader/dataset.py b/src/electrai/dataloader/dataset.py index a473c2da..d01996ad 100644 --- a/src/electrai/dataloader/dataset.py +++ b/src/electrai/dataloader/dataset.py @@ -79,10 +79,6 @@ def __init__( ] 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( @@ -120,10 +116,12 @@ def setup(self, stage: str): 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": + train_parts.append(AddDatasetID(splits["train"], dataset_id)) + val_parts.append(AddDatasetID(splits["validation"], dataset_id)) + elif stage == "test": + if "test" in splits and splits["test"] is not None: + test_parts.append(AddDatasetID(splits["test"], dataset_id)) if stage == "fit": self.train_set = ( diff --git a/src/electrai/lightning.py b/src/electrai/lightning.py index 560288fb..f56a9017 100644 --- a/src/electrai/lightning.py +++ b/src/electrai/lightning.py @@ -43,6 +43,7 @@ def validation_step(self, batch): return loss def _loss_calculation(self, batch): + # batch["Dataset_ID"] is available for future multi-head model extensions x = batch["data"] y = batch["label"] if isinstance(x, list): diff --git a/tests/electrai/dataloader/test_dataset.py b/tests/electrai/dataloader/test_dataset.py index 454edfa8..323e9dc9 100644 --- a/tests/electrai/dataloader/test_dataset.py +++ b/tests/electrai/dataloader/test_dataset.py @@ -1,19 +1,27 @@ from __future__ import annotations +import dataclasses import json from typing import TYPE_CHECKING +from unittest.mock import MagicMock, patch import numpy as np import pytest +import torch import zarr +from torch.utils.data import Dataset -from electrai.dataloader.dataset import RhoData +from electrai.dataloader.collate import collate_fn +from electrai.dataloader.dataset import AddDatasetID, DatasetSpec, RhoData, RhoRead from electrai.dataloader.utils import load_zarr if TYPE_CHECKING: from pathlib import Path +# --- Fixtures --- + + @pytest.fixture def rng(): return np.random.default_rng(seed=0) @@ -39,6 +47,24 @@ def filelist(tmp_path: Path, zarr_root: Path) -> Path: return filelist_path +class SimpleDataset(Dataset): + def __init__(self, n: int): + self.n = n + + def __len__(self): + return self.n + + def __getitem__(self, idx): + return { + "data": torch.tensor([float(idx)]), + "label": torch.tensor([float(idx)]), + "index": str(idx), + } + + +# --- TestLoadZarr --- + + class TestLoadZarr: def test_returns_arrays_divided_by_volume(self, zarr_root: Path): data, label = load_zarr(zarr_root, "mp-1") @@ -57,6 +83,9 @@ def test_missing_structure_attr_raises(self, tmp_path: Path, rng): load_zarr(tmp_path, "bad") +# --- TestRhoDataFormatDetection --- + + class TestRhoDataFormatDetection: def test_detects_zarr_format(self, filelist: Path): dataset = RhoData(str(filelist), precision="f32", augmentation=False) @@ -90,3 +119,270 @@ def test_getitem_zarr(self, filelist: Path): assert "label" in item assert "index" in item assert item["data"].shape[0] == 1 # unsqueeze adds channel dim + + +# --- TestAddDatasetID --- + + +class TestAddDatasetID: + def test_len_preserved(self): + base = SimpleDataset(10) + wrapped = AddDatasetID(base, functional_id=3) + assert len(wrapped) == 10 + + def test_dataset_id_added(self): + base = SimpleDataset(5) + wrapped = AddDatasetID(base, functional_id=7) + sample = wrapped[0] + assert sample["Dataset_ID"] == 7 + + def test_dataset_id_is_int(self): + base = SimpleDataset(3) + wrapped = AddDatasetID(base, functional_id="2") + assert isinstance(wrapped[0]["Dataset_ID"], int) + + def test_original_keys_preserved(self): + base = SimpleDataset(3) + wrapped = AddDatasetID(base, functional_id=1) + sample = wrapped[0] + assert "data" in sample + assert "label" in sample + assert "index" in sample + + def test_custom_key(self): + base = SimpleDataset(3) + wrapped = AddDatasetID(base, functional_id=5, key="src_id") + assert "src_id" in wrapped[0] + assert wrapped[0]["src_id"] == 5 + + +# --- TestCollateFn --- + + +class TestCollateFn: + def test_uniform_shapes_uses_default_collate(self): + batch = [ + { + "data": torch.zeros(4, 4, 4), + "label": torch.ones(4, 4, 4), + "index": "a", + "Dataset_ID": 0, + }, + { + "data": torch.zeros(4, 4, 4), + "label": torch.ones(4, 4, 4), + "index": "b", + "Dataset_ID": 0, + }, + ] + result = collate_fn(batch) + assert isinstance(result, dict) + assert result["data"].shape == (2, 4, 4, 4) + + def test_mismatched_shapes_fallback_returns_dict(self): + batch = [ + { + "data": torch.zeros(4, 4, 4), + "label": torch.ones(4, 4, 4), + "index": "a", + "Dataset_ID": 0, + }, + { + "data": torch.zeros(8, 8, 8), + "label": torch.ones(8, 8, 8), + "index": "b", + "Dataset_ID": 1, + }, + ] + result = collate_fn(batch) + assert isinstance(result, dict) + assert len(result["data"]) == 2 + assert len(result["label"]) == 2 + + def test_fallback_preserves_all_keys(self): + batch = [ + { + "data": torch.zeros(4, 4, 4), + "label": torch.ones(4, 4, 4), + "index": "a", + "Dataset_ID": 0, + }, + { + "data": torch.zeros(8, 8, 8), + "label": torch.ones(8, 8, 8), + "index": "b", + "Dataset_ID": 1, + }, + ] + result = collate_fn(batch) + assert set(result.keys()) == {"data", "label", "index", "Dataset_ID"} + + def test_fallback_correct_values(self): + batch = [ + { + "data": torch.zeros(2, 2, 2), + "label": torch.ones(2, 2, 2), + "index": "x", + "Dataset_ID": 3, + }, + { + "data": torch.zeros(3, 3, 3), + "label": torch.ones(3, 3, 3), + "index": "y", + "Dataset_ID": 5, + }, + ] + result = collate_fn(batch) + assert result["index"] == ["x", "y"] + assert result["Dataset_ID"] == [3, 5] + + +# --- TestDatasetSpec --- + + +class TestDatasetSpec: + def test_required_root(self): + spec = DatasetSpec(root="/some/path") + assert spec.root == "/some/path" + + def test_defaults_are_none(self): + spec = DatasetSpec(root="/p") + assert spec.split_file is None + assert spec.val_frac is None + assert spec.dataset_id is None + + def test_frozen(self): + spec = DatasetSpec(root="/p") + with pytest.raises(dataclasses.FrozenInstanceError): + spec.root = "/other" + + +# --- TestRhoReadInit --- + + +class TestRhoReadInit: + def test_empty_datasets_raises(self): + with pytest.raises(ValueError, match="at least one"): + RhoRead(datasets=[]) + + def test_none_datasets_raises(self): + with pytest.raises(ValueError, match="at least one"): + RhoRead(datasets=None) + + def test_auto_assigns_dataset_id(self): + reader = RhoRead( + datasets=[{"root": "/a/filelist.txt"}, {"root": "/b/filelist.txt"}] + ) + assert reader.specs[0].dataset_id == 0 + assert reader.specs[1].dataset_id == 1 + + def test_explicit_dataset_id_preserved(self): + reader = RhoRead(datasets=[{"root": "/a/filelist.txt", "dataset_id": 42}]) + assert reader.specs[0].dataset_id == 42 + + def test_default_val_frac_applied(self): + reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}], default_val_frac=0.1) + assert reader.specs[0].val_frac == 0.1 + + def test_per_spec_val_frac_overrides_default(self): + reader = RhoRead( + datasets=[{"root": "/a/filelist.txt", "val_frac": 0.2}], + default_val_frac=0.1, + ) + assert reader.specs[0].val_frac == 0.2 + + def test_default_split_file_applied(self): + reader = RhoRead( + datasets=[{"root": "/a/filelist.txt"}], + default_split_file="/default/split.json", + ) + assert reader.specs[0].split_file == "/default/split.json" + + def test_per_spec_split_file_overrides_default(self): + reader = RhoRead( + datasets=[{"root": "/a/filelist.txt", "split_file": "/custom/split.json"}], + default_split_file="/default/split.json", + ) + assert reader.specs[0].split_file == "/custom/split.json" + + +# --- TestRhoReadSetup --- + + +class TestRhoReadSetup: + def _make_fake_splits(self, n: int = 5): + ds = SimpleDataset(n) + from torch.utils.data import Subset + + train = Subset(ds, list(range(4))) + val = Subset(ds, [4]) + return {"train": train, "validation": val} + + @patch("electrai.dataloader.dataset.split_data") + @patch("electrai.dataloader.dataset.RhoData") + def test_fit_stage_builds_train_and_val(self, mock_rho, mock_split): + mock_rho.return_value = MagicMock(spec=Dataset) + mock_split.return_value = self._make_fake_splits() + + reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}]) + reader.setup("fit") + + assert reader.train_set is not None + assert reader.val_set is not None + assert reader.test_set is None + + @patch("electrai.dataloader.dataset.split_data") + @patch("electrai.dataloader.dataset.RhoData") + def test_test_stage_does_not_build_train_or_val(self, mock_rho, mock_split): + mock_rho.return_value = MagicMock(spec=Dataset) + splits = self._make_fake_splits() + splits["test"] = splits["train"] + mock_split.return_value = splits + + reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}]) + reader.setup("test") + + assert reader.test_set is not None + assert reader.train_set is None + assert reader.val_set is None + + @patch("electrai.dataloader.dataset.split_data") + @patch("electrai.dataloader.dataset.RhoData") + def test_multi_dataset_concat(self, mock_rho, mock_split): + from torch.utils.data import ConcatDataset + + mock_rho.return_value = MagicMock(spec=Dataset) + mock_split.return_value = self._make_fake_splits() + + reader = RhoRead( + datasets=[{"root": "/a/filelist.txt"}, {"root": "/b/filelist.txt"}] + ) + reader.setup("fit") + + assert isinstance(reader.train_set, ConcatDataset) + assert isinstance(reader.val_set, ConcatDataset) + + @patch("electrai.dataloader.dataset.split_data") + @patch("electrai.dataloader.dataset.RhoData") + def test_single_dataset_no_concat(self, mock_rho, mock_split): + from torch.utils.data import ConcatDataset + + mock_rho.return_value = MagicMock(spec=Dataset) + mock_split.return_value = self._make_fake_splits() + + reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}]) + reader.setup("fit") + + assert not isinstance(reader.train_set, ConcatDataset) + + @patch("electrai.dataloader.dataset.split_data") + @patch("electrai.dataloader.dataset.RhoData") + def test_dataset_id_propagated(self, mock_rho, mock_split): + mock_rho.return_value = MagicMock(spec=Dataset) + mock_split.return_value = self._make_fake_splits() + + reader = RhoRead(datasets=[{"root": "/a/filelist.txt", "dataset_id": 99}]) + reader.setup("fit") + + sample = reader.train_set[0] + assert sample["Dataset_ID"] == 99 From b48ae8e2feb2d07b0aee61be2be3b64fef861520 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Tue, 19 May 2026 17:22:44 -0400 Subject: [PATCH 3/6] Address second round of PR review comments - Handle list batch in test_step for mismatched-shape multi-dataset inference - Skip split_data for test stage when no split_file is defined (no test set possible) - Shallow-copy dict in AddDatasetID.__getitem__ to avoid mutating base dataset output Co-Authored-By: Claude Sonnet 4.6 --- src/electrai/dataloader/dataset.py | 5 ++++- src/electrai/lightning.py | 13 +++++++++++-- tests/electrai/dataloader/test_dataset.py | 4 +++- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/src/electrai/dataloader/dataset.py b/src/electrai/dataloader/dataset.py index d01996ad..253fd15e 100644 --- a/src/electrai/dataloader/dataset.py +++ b/src/electrai/dataloader/dataset.py @@ -35,7 +35,7 @@ def __len__(self): return len(self.base) def __getitem__(self, idx): - out = self.base[idx] + out = dict(self.base[idx]) out[self.key] = self.functional_id return out @@ -104,6 +104,9 @@ def setup(self, stage: str): test_parts: list[Dataset] = [] for spec in self.specs: + if stage == "test" and spec.split_file is None: + continue + ds = RhoData( spec.root, precision=self.precision, augmentation=self.augmentation ) diff --git a/src/electrai/lightning.py b/src/electrai/lightning.py index f56a9017..7aea6869 100644 --- a/src/electrai/lightning.py +++ b/src/electrai/lightning.py @@ -95,8 +95,17 @@ def test_step(self, batch): end = torch.cuda.Event(enable_timing=True) start.record() - preds = self(x) - loss = self.loss_fn(preds, y) + if isinstance(x, list): + preds_list, losses = [], [] + for x_i, y_i in zip(x, y, strict=True): + p = self(x_i.unsqueeze(0)) + preds_list.append(p) + losses.append(self.loss_fn(p, y_i.unsqueeze(0))) + preds = torch.cat(preds_list) + loss = torch.stack(losses).mean() + else: + preds = self(x) + loss = self.loss_fn(preds, y) end.record() torch.cuda.synchronize() diff --git a/tests/electrai/dataloader/test_dataset.py b/tests/electrai/dataloader/test_dataset.py index 323e9dc9..f2f12402 100644 --- a/tests/electrai/dataloader/test_dataset.py +++ b/tests/electrai/dataloader/test_dataset.py @@ -339,7 +339,9 @@ def test_test_stage_does_not_build_train_or_val(self, mock_rho, mock_split): splits["test"] = splits["train"] mock_split.return_value = splits - reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}]) + reader = RhoRead( + datasets=[{"root": "/a/filelist.txt", "split_file": "/split.json"}] + ) reader.setup("test") assert reader.test_set is not None From c330bfb486ca2c52bc5e258c95a38faceb90a3b4 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Wed, 20 May 2026 11:35:20 -0400 Subject: [PATCH 4/6] Address third round of PR review comments - Fix y.detach().cpu() crash in test_step when y is a list (fallback collate path) - Gate split_data call by stage to avoid computing unused train/val splits on test - Warn when a dataset is skipped in test stage due to missing split_file - Fix test.py: use cfg.precision instead of cfg.model_precision - Fix test.py: use devices=auto and infer num_nodes from env for multi-GPU testing Co-Authored-By: Claude Sonnet 4.6 --- src/electrai/dataloader/dataset.py | 26 +++++++++++++++++++------- src/electrai/entrypoints/test.py | 11 +++++++++-- src/electrai/lightning.py | 7 ++++++- 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/src/electrai/dataloader/dataset.py b/src/electrai/dataloader/dataset.py index 253fd15e..a4845f51 100644 --- a/src/electrai/dataloader/dataset.py +++ b/src/electrai/dataloader/dataset.py @@ -1,5 +1,6 @@ from __future__ import annotations +import warnings from dataclasses import dataclass from pathlib import Path from typing import Any @@ -105,24 +106,35 @@ def setup(self, stage: str): for spec in self.specs: if stage == "test" and spec.split_file is None: + warnings.warn( + f"Dataset with root '{spec.root}' has no split_file and will be " + "skipped for the test stage (test indices require a split_file).", + UserWarning, + stacklevel=2, + ) continue 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) if stage == "fit": + splits = split_data( + ds, + val_frac=float(spec.val_frac), + split_file=spec.split_file, + random_seed=self.random_seed, + ) train_parts.append(AddDatasetID(splits["train"], dataset_id)) val_parts.append(AddDatasetID(splits["validation"], dataset_id)) elif stage == "test": + splits = split_data( + ds, + val_frac=float(spec.val_frac), + split_file=spec.split_file, + random_seed=self.random_seed, + ) if "test" in splits and splits["test"] is not None: test_parts.append(AddDatasetID(splits["test"], dataset_id)) diff --git a/src/electrai/entrypoints/test.py b/src/electrai/entrypoints/test.py index 5b02ce25..20f55d7a 100644 --- a/src/electrai/entrypoints/test.py +++ b/src/electrai/entrypoints/test.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from pathlib import Path from types import SimpleNamespace @@ -48,12 +49,18 @@ def test(args): tmp_dir = log_dir / "tmp" for directory in [log_dir, tmp_dir]: directory.mkdir(exist_ok=True, parents=True) + local_world_size = int( + os.environ.get("LOCAL_WORLD_SIZE", torch.cuda.device_count()) + ) + world_size = int(os.environ.get("WORLD_SIZE", local_world_size)) + num_nodes = max(1, world_size // local_world_size) trainer = Trainer( logger=None, callbacks=None, accelerator="gpu" if torch.cuda.is_available() else "cpu", - devices=1, - precision=cfg.model_precision, + devices="auto", + num_nodes=num_nodes, + precision=cfg.precision, ) lit_model.test_cfg = SimpleNamespace( diff --git a/src/electrai/lightning.py b/src/electrai/lightning.py index 7aea6869..c1123880 100644 --- a/src/electrai/lightning.py +++ b/src/electrai/lightning.py @@ -113,8 +113,13 @@ def test_step(self, batch): self.log("test_loss", loss, prog_bar=True, sync_dist=True) + y_cpu = ( + torch.cat([t.unsqueeze(0) for t in y]).detach().cpu() + if isinstance(y, list) + else y.detach().cpu() + ) out = { - "target": y.detach().cpu(), + "target": y_cpu, "index": indices, "nmae": loss.detach().cpu(), "duration": elapsed, From 07a1cfbf7576b3bd1a19cfd317d573443f54f652 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Wed, 20 May 2026 11:46:04 -0400 Subject: [PATCH 5/6] Fix Dataset_ID type inconsistency in collate_fn fallback In the fallback path, per-key collation is now attempted with default_collate so scalar/uniform fields (e.g. Dataset_ID) become tensors matching the normal path, while only genuinely mismatched tensors remain as lists. Co-Authored-By: Claude Sonnet 4.6 --- src/electrai/dataloader/collate.py | 8 +++++++- tests/electrai/dataloader/test_dataset.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/electrai/dataloader/collate.py b/src/electrai/dataloader/collate.py index df5f31fb..426a04b0 100644 --- a/src/electrai/dataloader/collate.py +++ b/src/electrai/dataloader/collate.py @@ -7,4 +7,10 @@ def collate_fn(batch): try: return default_collate(batch) except RuntimeError: - return {k: [d[k] for d in batch] for k in batch[0]} + result = {} + for k in batch[0]: + try: + result[k] = default_collate([d[k] for d in batch]) + except RuntimeError: + result[k] = [d[k] for d in batch] + return result diff --git a/tests/electrai/dataloader/test_dataset.py b/tests/electrai/dataloader/test_dataset.py index f2f12402..5f9be729 100644 --- a/tests/electrai/dataloader/test_dataset.py +++ b/tests/electrai/dataloader/test_dataset.py @@ -234,7 +234,7 @@ def test_fallback_correct_values(self): ] result = collate_fn(batch) assert result["index"] == ["x", "y"] - assert result["Dataset_ID"] == [3, 5] + assert result["Dataset_ID"].tolist() == [3, 5] # --- TestDatasetSpec --- From 80fb7dc74ce83b88c09426476e81f73c6f799741 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Wed, 20 May 2026 12:27:27 -0400 Subject: [PATCH 6/6] Update e2e_train.py to use new RhoRead datasets API Co-Authored-By: Claude Sonnet 4.6 --- scripts/e2e_train.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/scripts/e2e_train.py b/scripts/e2e_train.py index db0879ab..385adbde 100755 --- a/scripts/e2e_train.py +++ b/scripts/e2e_train.py @@ -186,12 +186,11 @@ class Config: raise ValueError("no samples remain after filtering") datamodule = RhoRead( - root=str(filelist), + datasets=[{"root": str(filelist), "val_frac": 0.4}], precision="f32", batch_size=cfg.nbatch, train_workers=0, val_workers=0, - val_frac=0.4, augmentation=False, random_seed=seed, )