"""Lightning data module and split-index models for cell datasets."""
# torchcell/datamodules/cell.py
# [[torchcell.datamodules.cell]]
# https://github.com/Mjvolk3/torchcell/tree/main/torchcell/datamodules/cell.py
# Test file: torchcell/datamodules/test_cell.py
import hashlib
import json
import logging
import math
import os
import os.path as osp
import random
from collections import defaultdict
from collections.abc import Callable, Iterable, Mapping
from typing import Any, cast
import lightning as L
import pandas as pd
import torch
from pydantic import BaseModel, Field, model_validator
from torch.utils.data import DataLoader as TorchDataLoader
from torch_geometric.loader import DataLoader, PrefetchLoader
from torchcell.datamodels import ModelStrict
logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)
[docs]
class IndexSplit(ModelStrict):
"""A sorted list of indices together with its element count."""
indices: list[int] = Field(..., description="Must be sorted in ascending order")
count: int
[docs]
@model_validator(mode="before")
@classmethod
def check_sorted_indices(cls, values: dict[str, Any]) -> dict[str, Any]:
"""Validate that the indices list is sorted in ascending order."""
indices = values.get("indices")
if indices and not all(
indices[i] <= indices[i + 1] for i in range(len(indices) - 1)
):
raise ValueError("Indices must be sorted in ascending order")
return values
def __repr__(self) -> str:
"""Return a truncated representation showing the first few indices."""
max_indices = 3
indices_str = (
f"[{', '.join(map(str, self.indices[:max_indices]))}"
f"{', ...' if len(self.indices) > max_indices else ''}]"
)
return f"IndexSplit(indices={indices_str}, count={self.count})"
[docs]
class DatasetSplit(BaseModel):
"""Per-split index groupings keyed by phenotype, perturbation count, or dataset."""
phenotype_label_index: dict[str, IndexSplit] | None = None
perturbation_count_index: dict[int, IndexSplit] | None = None
dataset_name_index: dict[str, IndexSplit] | None = None
[docs]
class DataModuleIndexDetails(ModelStrict):
"""Detailed per-split index breakdown with summary reporting."""
methods: list[str]
train: DatasetSplit
val: DatasetSplit
test: DatasetSplit
[docs]
def df_summary(self) -> pd.DataFrame:
"""Return a DataFrame summarizing counts and ratios per split and index key."""
data: defaultdict[tuple[str, Any], defaultdict[str, int]] = defaultdict(
lambda: defaultdict(int)
)
totals: defaultdict[tuple[str, Any], int] = defaultdict(int)
for split in ["train", "val", "test"]:
split_data = getattr(self, split)
for index_type, index_data in split_data.dict().items():
if index_data is not None:
for key, index_split in index_data.items():
# Handle both IndexSplit objects and dictionaries
if isinstance(index_split, dict):
count = cast(int, index_split.get("count"))
else:
count = index_split.count
data[(index_type, key)][split] = count
totals[(index_type, key)] += count
summary_data = []
for (index_type, key), splits in data.items():
total = totals[(index_type, key)]
for split in ["train", "val", "test"]:
count = splits[split]
ratio = count / total if total > 0 else 0
summary_data.append(
{
"split": split,
"index_type": index_type,
"key": key,
"count": count,
"ratio": ratio,
"total": total,
}
)
# Explicit columns so a details object with no index data yields an EMPTY frame
# with its columns (and ``__str__`` its "(empty)" form) rather than a column-less
# frame on which ``df["split"]`` raises ``KeyError``.
df = pd.DataFrame(
summary_data,
columns=["split", "index_type", "key", "count", "ratio", "total"],
).astype({"count": int, "ratio": float, "total": int})
# Create a categorical column for 'split' with the desired order
df["split"] = pd.Categorical(
df["split"], categories=["train", "val", "test"], ordered=True
)
# Sort the DataFrame
df = df.sort_values(["split", "index_type", "key"])
df["ratio"] = df["ratio"].round(3)
df = df.reset_index(drop=True)
return df
def __str__(self) -> str:
"""Return the summary DataFrame rendered as a string."""
df = self.df_summary()
if df.empty:
return "DataModuleIndexDetails(empty)"
return df.to_string()
[docs]
class DataModuleIndex(ModelStrict):
"""Train/val/test index lists that are sorted and mutually non-overlapping."""
train: list[int] = Field(..., description="Must be sorted in ascending order")
val: list[int] = Field(..., description="Must be sorted in ascending order")
test: list[int] = Field(..., description="Must be sorted in ascending order")
[docs]
@model_validator(mode="before")
@classmethod
def check_sorted_and_unique_indices(cls, values: dict[str, Any]) -> dict[str, Any]:
"""Validate each split is sorted and the splits do not overlap."""
for split in ["train", "val", "test"]:
indices = values.get(split, [])
if not all(indices[i] <= indices[i + 1] for i in range(len(indices) - 1)):
raise ValueError(f"{split} indices must be sorted in ascending order")
all_indices = (
values.get("train", []) + values.get("val", []) + values.get("test", [])
)
if len(set(all_indices)) != len(all_indices):
raise ValueError("Indices in train, val, and test must not overlap")
return values
def __repr__(self) -> str:
"""Return a truncated representation of the train/val/test indices."""
max_indices = 3
train_str_index = f"[{', '.join(map(str, self.train[:max_indices]))}{', ...' if len(self.train) > max_indices else ''}]"
val_str_index = f"[{', '.join(map(str, self.val[:max_indices]))}{', ...' if len(self.val) > max_indices else ''}]"
test_str_index = f"[{', '.join(map(str, self.test[:max_indices]))}{', ...' if len(self.test) > max_indices else ''}]"
return f"DataModuleIndex(train={train_str_index}, val={val_str_index}, test={test_str_index})"
def __str__(self) -> str:
"""Return a truncated representation with per-split index counts."""
max_indices = 3
train_str_index = f"[{', '.join(map(str, self.train[:max_indices]))}{', ...' if len(self.train) > max_indices else ''}]"
val_str_index = f"[{', '.join(map(str, self.val[:max_indices]))}{', ...' if len(self.val) > max_indices else ''}]"
test_str_index = f"[{', '.join(map(str, self.test[:max_indices]))}{', ...' if len(self.test) > max_indices else ''}]"
train_str = train_str_index + f" ({len(self.train)} indices)"
val_str = val_str_index + f" ({len(self.val)} indices)"
test_str = test_str_index + f" ({len(self.test)} indices)"
return f"DataModuleIndex(train={train_str}, val={val_str}, test={test_str})"
class DatasetIndexSplit(ModelStrict):
"""Per-dataset index lists grouped by train/val/test split."""
train: dict[str | int, list[int]]
val: dict[str | int, list[int]]
test: dict[str | int, list[int]]
def overlap_dataset_index_split(
dataset_index: dict[str | int, list[int]], data_module_index: DataModuleIndex
) -> DatasetIndexSplit:
"""Intersect each dataset's indices with the train/val/test split indices.
A split that holds none of the datasets is an EMPTY dict, not ``None``: an empty
overlap is a legitimate outcome (e.g. a small subset whose test split draws from one
dataset only), and the experiment 003 plotting scripts that call this treat an empty
split as "nothing to draw".
"""
train_set = set(data_module_index.train)
val_set = set(data_module_index.val)
test_set = set(data_module_index.test)
train_dict: dict[str | int, list[int]] = {}
val_dict: dict[str | int, list[int]] = {}
test_dict: dict[str | int, list[int]] = {}
for dataset_name, indices in dataset_index.items():
train_indices = sorted(list(set(indices) & train_set))
val_indices = sorted(list(set(indices) & val_set))
test_indices = sorted(list(set(indices) & test_set))
if train_indices:
train_dict[dataset_name] = train_indices
if val_indices:
val_dict[dataset_name] = val_indices
if test_indices:
test_dict[dataset_name] = test_indices
return DatasetIndexSplit(train=train_dict, val=val_dict, test=test_dict)
WORKER_TIMEOUT_S = 10800
def worker_dataloader_kwargs(
*,
batch_size: int,
shuffle: bool,
num_workers: int,
persistent_workers: bool,
pin_memory: bool,
prefetch_factor: int,
) -> dict[str, Any]:
"""Return the ``DataLoader`` keyword arguments every torchcell datamodule shares.
``CellDataModule`` and ``PerturbationSubsetDataModule`` both build their loaders
from this one dict (issue #580), so the two cannot drift apart again. The worker
options are guarded on ``num_workers``: with workers, persistence is as requested,
``timeout`` is ``WORKER_TIMEOUT_S`` (3 h), the context is ``spawn`` and
``prefetch_factor`` is passed through; at ``num_workers=0`` persistence is off,
``timeout`` is 0, the context is ``None`` and ``prefetch_factor`` is ``None``.
"""
return {
"batch_size": batch_size,
"shuffle": shuffle,
"num_workers": num_workers,
"persistent_workers": persistent_workers if num_workers > 0 else False,
"pin_memory": pin_memory,
# THE SAME GUARD AS `prefetch_factor` BELOW, and for the same reason -- this one
# was missed when that fix landed (commit 2cca6d83), so num_workers=0 was still
# unconstructible afterwards. `timeout` is a WORKER-QUEUE timeout, so at
# num_workers=0 torch asserts outright:
# AssertionError: _SingleProcessDataLoaderIter requires timeout == 0
# raised from `iter(dataloader)`, i.e. during sanity-check, AFTER the dataset and
# embeddings have been loaded. A trial therefore burns its full startup cost and
# then dies, which is why 119 of them could fail in a job that still exited
# COMPLETED 0:0.
"timeout": (WORKER_TIMEOUT_S if num_workers > 0 else 0),
"multiprocessing_context": ("spawn" if num_workers > 0 else None),
# Torch REJECTS a non-None prefetch_factor at num_workers=0 ("could only be
# specified in multiprocessing"), so this needs the same guard its two
# neighbors already carry. Without it num_workers=0 is not merely slow, it is
# unconstructible -- every DataLoader raises ValueError before the first batch.
#
# That is not hypothetical: on 2026-07-28 Delta job 20556837 ran NUM_WORKERS=0
# and all 119 optuna trials died here, ~15s apiece, while the job still exited
# COMPLETED 0:0. num_workers=0 is worth supporting because on a parallel
# filesystem the spawn path is the expensive one -- each worker re-imports the
# whole stack off /work/hdd.
"prefetch_factor": (prefetch_factor if num_workers > 0 else None),
}
[docs]
class CellDataModule(L.LightningDataModule):
"""Lightning data module that splits a cell dataset and builds dataloaders."""
def __init__(
self,
dataset: Any, # dynamic cell dataset duck-typed by split-index attrs
cache_dir: str = "cache",
batch_size: int = 32,
random_seed: int = 42,
num_workers: int = 0,
pin_memory: bool = False,
prefetch: bool = False,
prefetch_factor: int = 2,
persistent_workers: bool = True,
split_indices: str | list[str] | None = None,
follow_batch: list[str] | None = None,
train_shuffle: bool = True,
collate_fn: Callable[[list[Any]], Any] | None = None,
val_batch_size: int | None = None,
pinned_test_indices: Iterable[int] | None = None,
pinned_split_indices: Mapping[str, Iterable[int]] | None = None,
index_subset: Iterable[int] | None = None,
) -> None:
"""Store dataloader/split configuration and compute the split indices.
``pinned_test_indices`` forces those record indices into the TEST split, overriding
the random assignment. It reproduces an EXTERNAL split inside ours -- e.g. Merzbacher
2025's betaxanthin test ORFs, so their published numbers and ours are computed on the
same genes. The remaining records are split train/val/test by ``random_seed`` as
usual, so sweeping the seed still re-rolls train/val while the pinned comparison set
stays fixed.
``pinned_split_indices`` generalizes that to ALL THREE splits: a mapping like
``{"train": [...], "val": [...], "test": [...]}`` whose records are forced into
the named split. It reproduces a PRIOR experiment's full assignment inside a
larger dataset -- e.g. 010's tmi train/val/test carried into 025 by genotype
identity -- so a comparison across datasets is computed on identical splits of
the shared records while everything new is seed-assigned around them. A record
may appear in at most one pinned split; overlap raises at construction.
``index_subset`` restricts the RECORD POOL before any splitting, so a record
outside it lands in no split at all. Pinning assigns but never excludes, which is
why it cannot express a subset on its own: 025 holds 13,525,071 records of which
376,732 are the trigenic triples, and an arm training on triples only needs the
other 13.1M gone, not placed somewhere. The subset is intersected into each split
index key BEFORE shuffling rather than filtered out afterwards, so an arm over
376k records costs 376k of work on a 13.5M-record dataset.
Both the pinned splits and the subset are hashed into the split-cache filename, so
changing either selects a different cache file rather than silently loading the
previous arm's index.
"""
super().__init__()
self.dataset = dataset
self.cache_dir = cache_dir
self.batch_size = batch_size
self.val_batch_size = (
val_batch_size if val_batch_size is not None else batch_size
)
self.random_seed = random_seed
self.num_workers = num_workers
self.pin_memory = pin_memory
self.prefetch = prefetch
self.prefetch_factor = prefetch_factor
self.persistent_workers = persistent_workers
self.train_shuffle = train_shuffle
self.train_ratio = 0.8
self.val_ratio = 0.1
self.pinned_test_indices: set[int] = (
set(pinned_test_indices) if pinned_test_indices is not None else set()
)
self.index_subset: set[int] | None = (
set(index_subset) if index_subset is not None else None
)
assert self.index_subset is None or self.index_subset, (
"index_subset is empty; pass None to use every record"
)
self.pinned_split_indices: dict[str, set[int]] = {
split: set(indices)
for split, indices in (pinned_split_indices or {}).items()
}
unknown = set(self.pinned_split_indices) - {"train", "val", "test"}
assert not unknown, f"pinned_split_indices has unknown splits: {unknown}"
pinned_sets = list(self.pinned_split_indices.values())
for i, a in enumerate(pinned_sets):
for b in pinned_sets[i + 1 :]:
overlap = a & b
assert not overlap, (
f"pinned_split_indices overlap across splits: {sorted(overlap)[:5]}..."
)
self.split_indices = (
split_indices
if isinstance(split_indices, list)
else [split_indices]
if split_indices
else []
)
# The split is defined BY these keys: a record is placed only through the keys
# that hold it, so with none every record lands in no split and the module would
# train on nothing without an error. Refused here, before the index is computed.
if not self.split_indices:
raise ValueError(
"split_indices is required: name at least one split index of the dataset "
"(e.g. 'phenotype_label_index'); without one every record lands in no split"
)
self._index: DataModuleIndex | None = None
self._index_details: DataModuleIndexDetails | None = None
if follow_batch is None:
self.follow_batch = ["x", "x_pert"]
else:
self.follow_batch = follow_batch
self.collate_fn = collate_fn
# Compute index during initialization
self.index
self.index_details
@property
def index(self) -> DataModuleIndex:
"""Return the train/val/test split index, computing or loading it if needed."""
if self._index is None or not self._cached_files_exist():
self._load_or_compute_index()
return self._index # type: ignore[return-value] # _load_or_compute_index() always populates _index
@property
def index_details(self) -> DataModuleIndexDetails:
"""Return the detailed split-index breakdown, computing or loading if needed."""
if self._index_details is None or not self._cached_files_exist():
self._load_or_compute_index()
return self._index_details # type: ignore[return-value] # _load_or_compute_index() always populates _index_details
def _pin_tag(self) -> str:
"""Cache-key suffix identifying the pinned test set, empty when there is none.
The cached index is keyed by seed alone, so WITHOUT this a pinned run would silently
load an UNPINNED `index_seed_42.json` left by an earlier run and train on the wrong
split -- with no error, and with the pinned genes back in train. The tag hashes the
pinned index set so any change to it (or removing the pin) selects a different file.
"""
if not self.pinned_test_indices and not self.pinned_split_indices:
return ""
parts = [",".join(map(str, sorted(self.pinned_test_indices)))]
n_pinned = len(self.pinned_test_indices)
for split in ("train", "val", "test"):
indices = self.pinned_split_indices.get(split, set())
parts.append(f"{split}:" + ",".join(map(str, sorted(indices))))
n_pinned += len(indices)
payload = "|".join(parts).encode()
return f"_pin{n_pinned}-{hashlib.sha256(payload).hexdigest()[:8]}"
def _subset_tag(self) -> str:
"""Cache-key suffix identifying the record subset, empty when there is none.
Separate from ``_pin_tag`` because the two vary independently: the 025 ladder holds
the split fixed (R or Q) while the training pool changes S0 -> S2 -> S3, and holds
the pool fixed while the split regime changes. Without this an S2 run would load
S0's cached index and train on 376,732 records while reporting 382,426.
"""
if self.index_subset is None:
return ""
payload = ",".join(map(str, sorted(self.index_subset))).encode()
return f"_sub{len(self.index_subset)}-{hashlib.sha256(payload).hexdigest()[:8]}"
def _cache_files(self) -> tuple[str, str]:
"""Paths of the cached index and index-details files for this configuration."""
tag = f"{self._pin_tag()}{self._subset_tag()}"
return (
osp.join(self.cache_dir, f"index_seed_{self.random_seed}{tag}.json"),
osp.join(
self.cache_dir, f"index_details_seed_{self.random_seed}{tag}.json"
),
)
def _restrict(self, indices: Iterable[int]) -> list[int]:
"""Drop indices outside ``index_subset``; a no-op when no subset is configured."""
if self.index_subset is None:
return list(indices)
return [i for i in indices if i in self.index_subset]
def _load_or_compute_index(self) -> None:
os.makedirs(self.cache_dir, exist_ok=True)
index_file, details_file = self._cache_files()
if osp.exists(index_file) and osp.exists(details_file):
try:
with open(index_file) as f:
log.info(f"Loading index from {index_file}")
index_dict = json.load(f)
self._index = DataModuleIndex(**index_dict)
with open(details_file) as f:
log.info(f"Loading index details from {details_file}")
details_dict = json.load(f)
self._index_details = DataModuleIndexDetails(**details_dict)
except Exception as e:
print(f"Error loading index or details: {e}. Regenerating...")
self._compute_and_save_index(index_file, details_file)
else:
self._compute_and_save_index(index_file, details_file)
def _compute_and_save_index(self, index_file: str, details_file: str) -> None:
log.info("Generating detailed index...")
random.seed(self.random_seed)
all_indices = set(range(len(self.dataset)))
if self.index_subset is not None:
outside = len(self.index_subset - all_indices)
assert not outside, (
f"index_subset names {outside} indices outside the "
f"{len(self.dataset)}-record dataset"
)
all_indices = set(self.index_subset)
log.info(
f"index_subset restricts the pool to {len(all_indices)} of "
f"{len(self.dataset)} records before splitting."
)
split_data: dict[str, defaultdict[str, set[int]]] = {
"train": defaultdict(set),
"val": defaultdict(set),
"test": defaultdict(set),
}
# First, split each index independently
for index_name in self.split_indices:
original_index = getattr(self.dataset, index_name)
for key, indices in original_index.items():
# Narrow to the subset HERE, not after assignment. Filtering afterwards
# would still shuffle and ratio-split all 13.5M records of the 025 build
# to keep 376k of them, and would draw the 80/10/10 boundaries on the full
# key rather than on the part of it the arm actually trains on.
indices = self._restrict(indices)
if not indices:
continue
random.shuffle(indices)
num_samples = len(indices)
num_train = int(self.train_ratio * num_samples)
num_val = int(self.val_ratio * num_samples)
split_data["train"][index_name].update(indices[:num_train])
split_data["val"][index_name].update(
indices[num_train : num_train + num_val]
)
split_data["test"][index_name].update(indices[num_train + num_val :])
# Then, create initial final splits
final_splits = {
"train": all_indices.intersection(
*[split_data["train"][index] for index in self.split_indices]
),
"val": all_indices.intersection(
*[split_data["val"][index] for index in self.split_indices]
),
"test": all_indices.intersection(
*[split_data["test"][index] for index in self.split_indices]
),
}
# A record may appear under SEVERAL keys of the same split index, and those keys
# are shuffled independently -- so it can be drawn into train under one key and
# into val under another, landing in two of the intersections above. That is a
# train/val/test leak, and `DataModuleIndex` rightly refuses it.
#
# It is not hypothetical: under deletion-keyed aggregation one genotype group
# holds records with DIFFERENT perturbation counts (a betaxanthin record carries
# its 4-gene cassette + the deletion = 5, a beta-carotene record 4, a metabolome
# record 1), so `perturbation_count_index` files that single record under three
# keys. The same happens in `phenotype_label_index` for any multi-modality
# genotype.
#
# Conflicts are pulled OUT of every split and returned to `remaining`, so the
# ratio-balancing assignment below places each conflicted record exactly once.
# When no record is multiply-keyed this is a no-op and the split is unchanged.
conflicted = (
(final_splits["train"] & final_splits["val"])
| (final_splits["train"] & final_splits["test"])
| (final_splits["val"] & final_splits["test"])
)
if conflicted:
log.info(
f"{len(conflicted)} records were assigned to more than one split by "
"independently-shuffled index keys; reassigning them exactly once."
)
for split in ("train", "val", "test"):
final_splits[split] -= conflicted
# Sophisticated assignment of remaining indices
remaining = all_indices - (
final_splits["train"] | final_splits["val"] | final_splits["test"]
)
for index_name in self.split_indices:
original_index = getattr(self.dataset, index_name)
for key, indices in original_index.items():
key_indices = set(self._restrict(indices))
key_remaining = key_indices & remaining
if not key_remaining:
continue
current_counts = {
split: len(key_indices & final_splits[split])
for split in ["train", "val", "test"]
}
total_count = sum(current_counts.values()) + len(key_remaining)
# Targets partition the key exactly: test is the REMAINDER, as in the
# first-pass slicing, not ``int(n * (1 - train - val))``, whose ratio is
# 0.09999999999999995 in floating point and so gave a 10-record key a test
# target of 0 (and a ZeroDivisionError below).
target_train = int(self.train_ratio * total_count)
target_val = int(self.val_ratio * total_count)
target_counts = {
"train": target_train,
"val": target_val,
"test": total_count - target_train - target_val,
}
empty_targets = [
split for split, target in target_counts.items() if target == 0
]
if empty_targets:
raise ValueError(
f"Key {key!r} of {index_name} has {len(key_remaining)} record(s) "
f"left to balance, but its {total_count} "
f"records give a target of 0 for {empty_targets}; balancing needs "
f"at least {math.ceil(1 / self.val_ratio)} records per such key"
)
for idx in key_remaining:
best_split = min(
["train", "val", "test"],
key=lambda x: (
(current_counts[x] - target_counts[x]) / target_counts[x]
),
)
final_splits[best_split].add(idx)
current_counts[best_split] += 1
remaining.remove(idx)
# PINNED TEST SET, applied last so it overrides every ratio-driven assignment above.
# Reproduces an EXTERNAL split inside ours: the pinned records go to test and are
# removed from train/val, and everything else keeps its seed-driven assignment. That
# is what makes a head-to-head comparison possible -- their metric and ours are then
# computed on the same genes -- while sweeping `random_seed` still re-rolls train/val
# around a comparison set that never moves.
#
# Test therefore ends up LARGER than `1 - train_ratio - val_ratio`, by design; the
# ratios govern the remaining pool, not the pinned block.
if self.pinned_test_indices:
pinned = self.pinned_test_indices & all_indices
missing = len(self.pinned_test_indices) - len(pinned)
final_splits["train"] -= pinned
final_splits["val"] -= pinned
final_splits["test"] |= pinned
log.info(
f"Pinned {len(pinned)} records into test "
f"({missing} requested indices are outside the record pool). "
f"Splits: train={len(final_splits['train'])} "
f"val={len(final_splits['val'])} test={len(final_splits['test'])}"
)
# FULL PINNED SPLITS, applied last of all so they override both the ratio-driven
# assignment and the test-only pin above. Each pinned record is removed from every
# split and placed exactly where the mapping says, reproducing a prior experiment's
# complete train/val/test assignment (e.g. 010's tmi splits inside 025) while the
# unpinned remainder keeps its seed-driven assignment.
if self.pinned_split_indices:
for split_name, requested in self.pinned_split_indices.items():
pinned = requested & all_indices
for other in ("train", "val", "test"):
final_splits[other] -= pinned
final_splits[split_name] |= pinned
log.info(
f"Pinned {len(pinned)} records into {split_name} "
f"({len(requested) - len(pinned)} requested indices are outside "
"the record pool)."
)
log.info(
f"Splits after full pinning: train={len(final_splits['train'])} "
f"val={len(final_splits['val'])} test={len(final_splits['test'])}"
)
# Every record of the pool must be placed. A record under no split-index key is
# never reached by the key-driven assignment above, and dropping it silently
# would be a data loss.
orphans = sorted(
all_indices
- (final_splits["train"] | final_splits["val"] | final_splits["test"])
)
if orphans:
shown = ", ".join(map(str, orphans[:10]))
more = f", ... ({len(orphans)} in all)" if len(orphans) > 10 else ""
raise ValueError(
f"{len(orphans)} record(s) sit under no key of {self.split_indices} and "
f"would land in no split: [{shown}{more}]"
)
# Create DataModuleIndexDetails object
self._index_details = DataModuleIndexDetails(
methods=self.split_indices,
train=DatasetSplit(),
val=DatasetSplit(),
test=DatasetSplit(),
)
for split in ["train", "val", "test"]:
for index_name in self.split_indices:
original_index = getattr(self.dataset, index_name)
split_index_data: dict[str | int, IndexSplit] = {}
for key, indices in original_index.items():
intersect = sorted(
set(self._restrict(indices)) & final_splits[split]
)
split_index_data[key] = IndexSplit(
indices=intersect, count=len(intersect)
)
setattr(
getattr(self._index_details, split), index_name, split_index_data
)
# Create DataModuleIndex object
self._index = DataModuleIndex(
train=sorted(list(final_splits["train"])),
val=sorted(list(final_splits["val"])),
test=sorted(list(final_splits["test"])),
)
# Save the index and details separately
with open(index_file, "w") as f:
json.dump(self._index.model_dump(), f, indent=2)
with open(details_file, "w") as f:
json.dump(self._index_details.model_dump(), f, indent=2)
def _cached_files_exist(self) -> bool:
index_file, details_file = self._cache_files()
return osp.exists(index_file) and osp.exists(details_file)
[docs]
def setup(self, stage: str | None = None) -> None:
"""Build train/val/test Subset datasets from the computed split indices.
IDEMPOTENT BY DESIGN. Lightning calls ``setup()`` again inside ``trainer.fit()``
even when the caller already invoked it, so an unconditional rebuild silently
discarded any post-setup narrowing of ``Subset.indices`` (e.g. the
``require_modalities`` filter in experiments/019, which logged "4074 -> 1161 rows"
while every epoch still ran 128 batches = 4074/32). Rebuilding only when the
subsets do not yet exist makes the caller's filter survive the second call.
"""
if (
getattr(self, "train_dataset", None) is not None
and getattr(self, "val_dataset", None) is not None
and getattr(self, "test_dataset", None) is not None
):
return
self.train_dataset = torch.utils.data.Subset(self.dataset, self.index.train)
self.val_dataset = torch.utils.data.Subset(self.dataset, self.index.val)
self.test_dataset = torch.utils.data.Subset(self.dataset, self.index.test)
def _get_dataloader(
self,
dataset: Any, # dynamic dataset/Subset passed through to the loaders
shuffle: bool = False,
batch_size: int | None = None,
) -> TorchDataLoader[Any] | PrefetchLoader:
"""Build the loader for ``dataset``: PyG's by default, torch's for a custom collate.
PyG's ``DataLoader`` pops any ``collate_fn`` and installs its own ``Collater``, so
a caller's function (e.g. experiment 006's ``LazyCollater``) would never run. When
``collate_fn`` is given the plain torch ``DataLoader`` is used with it instead, and
``follow_batch`` is not applied: that collate owns the batching.
"""
# Use provided batch_size or fall back to self.batch_size
if batch_size is None:
batch_size = self.batch_size
dataloader_kwargs = worker_dataloader_kwargs(
batch_size=batch_size,
shuffle=shuffle,
num_workers=self.num_workers,
persistent_workers=self.persistent_workers,
pin_memory=self.pin_memory,
prefetch_factor=self.prefetch_factor,
)
loader: TorchDataLoader[Any]
if self.collate_fn is not None:
loader = TorchDataLoader(
dataset, collate_fn=self.collate_fn, **dataloader_kwargs
)
else:
loader = DataLoader(
dataset, follow_batch=self.follow_batch, **dataloader_kwargs
)
if self.prefetch:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
return PrefetchLoader(loader, device=device)
return loader
[docs]
def train_dataloader(self) -> TorchDataLoader[Any] | PrefetchLoader:
"""Return a dataloader over the training subset."""
return self._get_dataloader(self.train_dataset, shuffle=self.train_shuffle)
[docs]
def val_dataloader(self) -> TorchDataLoader[Any] | PrefetchLoader:
"""Return a dataloader over the validation subset."""
return self._get_dataloader(self.val_dataset, batch_size=self.val_batch_size)
[docs]
def test_dataloader(self) -> TorchDataLoader[Any] | PrefetchLoader:
"""Return a dataloader over the test subset."""
return self._get_dataloader(self.test_dataset)
[docs]
def all_dataloader(self) -> TorchDataLoader[Any] | PrefetchLoader:
"""Return a dataloader over the entire dataset."""
return self._get_dataloader(self.dataset)
if __name__ == "__main__":
pass