diff --git a/scripts/e2e_train.py b/scripts/e2e_train.py index db0879a..385adbd 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, ) diff --git a/src/electrai/configs/MP/config_resnet.yaml b/src/electrai/configs/MP/config_resnet.yaml index 482c48a..244444f 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 # ignored when split_file is provided 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 b270db2..3825d8c 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 # ignored when split_file is provided 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 144924f..426a04b 100644 --- a/src/electrai/dataloader/collate.py +++ b/src/electrai/dataloader/collate.py @@ -7,5 +7,10 @@ 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) + 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/src/electrai/dataloader/dataset.py b/src/electrai/dataloader/dataset.py index 0a0a1c1..a4845f5 100644 --- a/src/electrai/dataloader/dataset.py +++ b/src/electrai/dataloader/dataset.py @@ -1,66 +1,154 @@ from __future__ import annotations +import warnings +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 = dict(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): + 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: + 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 + ) + 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)) + 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( diff --git a/src/electrai/entrypoints/test.py b/src/electrai/entrypoints/test.py index 5b02ce2..20f55d7 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 560288f..c112388 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): @@ -94,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() @@ -103,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, diff --git a/tests/electrai/dataloader/test_dataset.py b/tests/electrai/dataloader/test_dataset.py index 454edfa..5f9be72 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,272 @@ 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"].tolist() == [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", "split_file": "/split.json"}] + ) + 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