Support the combination of multiple datasets with the dataloader

[hanaol]

Problem
The current dataloader only supports a single dataset. As a result, it is not possible to combine data coming from different paths.

Solution
This PR extends the dataloader to support multiple datasets. The root, split_file, and val_frac parameters retain the same meaning as before. In addition, this introduces a new parameter, dataset_id, which serves as an identifier for each dataset and will be useful for future development of multi-head models.

Example

  datasets:
    - root: $DATASET1_PATH/mp_filelist.txt
      split_file: $DATASET1_PATH/split.json
      val_frac: 0.005
    - root: $DATASET2_PATH/mp_filelist.txt
      split_file: $DATASET2_PATH/split.json
      val_frac: 0.005

[forklady42]

@claude please review this PR

(I'll review afterwards)

[claude]

Claude finished @forklady42's task in 1m 58s —— View job

---

PR Review: Support the combination of multiple datasets with the dataloader

• Gather context
• Read changed files
• Analyze code
• Post review

---

Overall this is a well-structured approach to multi-dataset support. The DatasetSpec dataclass, AddDatasetID wrapper, and ConcatDataset usage are all clean patterns. A few issues worth addressing before merging:

Bug: collate_fn fallback is incorrect for dict-based batches

collate.py:10-11 — The fallback triggers when default_collate raises a RuntimeError (e.g. mismatched tensor shapes). At that point, batch is a list of dicts like {"data": ..., "label": ..., "index": ..., "Dataset_ID": ...}. But zip(*batch) unpacks dict keys, not values — so x, y, index, dataset_id would receive ("data","data",...), ("label","label",...), etc. rather than actual tensors.

This was a pre-existing issue (the old fallback had the same structure when RhoData.__getitem__ still returned a dict), but the new dataset_id field makes it more prominent. The fallback path also returns a positional 4-tuple while the normal default_collate path returns a dict — inconsistent return types will break any downstream consumer.

Fix this →

Performance: setup() always processes all datasets regardless of stage

dataset.py:105-138 — setup() iterates all specs and calls split_data() on every dataset unconditionally, even when stage == "test" (in which case train_parts and val_parts are built and then discarded). For large multi-dataset configurations this loads and splits unnecessary data.

# Current: all parts built regardless of stage
for spec in self.specs:
    ...
    train_parts.append(...)   # wasted when stage == "test"
    val_parts.append(...)     # wasted when stage == "test"

Consider gating which parts are built based on stage.

Fix this →

Design note: dataset_id is silently ignored by the model

The "Dataset_ID" key is added to every batch sample (dataset.py:39) but lightning.py:_loss_calculation only accesses batch["data"] and batch["label"]. For training and validation, the dataset_id is a no-op today. The PR description flags this as intentional ("useful for future development of multi-head models"), so this is fine — but worth a comment in _loss_calculation or the docstring so future developers know to extend it there.

Minor: defensive None check on a non-optional field

dataset.py:82 — DatasetSpec.root is typed as str (no default, required). The if s.root is None guard is dead code under normal construction. Either change the type to str | None = None (to make it optional), or remove the guard and let the dataclass raise a TypeError naturally.

Missing tests

No tests were added for the new multi-dataset path (DatasetSpec, AddDatasetID, multi-spec setup()). Given the collate fallback bug above, a test covering mismatched-shape batches across datasets would be valuable.

---

[forklady42]

Minor comment about clarifying the config interface. Also worth reviewing Claude's comments.

[forklady42]

val_frac isn't used when split_file is provided. this is at least worth a comment so no one is misled

[hanaol]

@claude review this PR.

[claude]

Claude finished @hanaol's task in 1m 43s —— View job

---

PR Review: Support the combination of multiple datasets with the dataloader

• Gather context
• Read changed files
• Analyze code
• Post review

---

This revision addresses the main issues from the prior review. Here's the current state:

What was fixed

• collate_fn fallback (collate.py:10): Correctly fixed to {k: [d[k] for d in batch] for k in batch[0]} — returns a dict and properly extracts values rather than keys. The tests in TestCollateFn cover both the uniform and mismatched-shape paths well.
• dataset_id no-op comment: lightning.py:46 now has # batch["Dataset_ID"] is available for future multi-head model extensions. Clear and sufficient.
• Config val_frac comment: Both configs now have # ignored when split_file is provided, addressing forklady42's comment.
• Tests: Comprehensive suite added — TestAddDatasetID, TestCollateFn, TestDatasetSpec, TestRhoReadInit, and TestRhoReadSetup all look solid.

Remaining issues

Bug: test_step doesn't handle the list (fallback collate) path

lightning.py:91-98 — _loss_calculation handles both tensor and list batches, but test_step does not:

def test_step(self, batch):
    x = batch["data"]   # could be a list if collate fallback triggered
    y = batch["label"]
    ...
    preds = self(x)     # will fail if x is a list
    loss = self.loss_fn(preds, y)

For multi-dataset test runs where samples have mismatched shapes, the fallback collate returns lists — and self(x) will fail since the model expects a tensor. _loss_calculation already has the right loop pattern; test_step should use a similar approach or delegate to it.

Fix this →

Performance: split_data still called for all specs regardless of stage

dataset.py:106-115 — RhoData and split_data are called unconditionally for every spec before the if stage == "fit" / elif stage == "test" branching. When stage == "test", train and val splits are computed and immediately discarded. If split_data is expensive (reading file indices, shuffling), this is wasteful.

for spec in self.specs:
    ds = RhoData(...)      # always called
    splits = split_data(...)  # always called — train+val splits computed even for stage="test"
    
    if stage == "fit":
        ...
    elif stage == "test":
        ...

This was flagged in the previous review and is still present. Given the stage flag is available, the split logic could be gated.

Minor: AddDatasetID.__getitem__ mutates the base dataset's returned dict

dataset.py:39 — out[self.key] = self.functional_id modifies the dict in-place. RhoData returns a fresh dict each call so this is safe right now, but it's a subtle footgun if AddDatasetID is wrapped around a dataset that caches or shares its return value. A out = dict(self.base[idx]) shallow copy would make this safe in general.

---

Overall the PR is in good shape — the multi-dataset design is clean, tests are solid, and the major bugs from the previous round were fixed. The test_step list-handling gap is the only functional bug remaining.

[hanaol]

@claude review this PR

[claude]

Claude finished @hanaol's task in 1m 23s —— View job

---

PR Review: Support the combination of multiple datasets with the dataloader

• Gather context
• Read changed files
• Analyze code
• Post review

---

This revision addresses all of the major issues from prior rounds. Here's the current state:

What was fixed since the last review

• test_step list-handling (lightning.py:98-108): Now correctly handles the fallback-collate list path with a per-sample loop — same pattern as _loss_calculation. Good.
• AddDatasetID mutation (dataset.py:38): out = dict(self.base[idx]) creates a shallow copy, making the in-place out[self.key] = ... safe. Good.
• collate_fn fallback: Dict-based, correct.
• Config comments (config_resnet.yaml:7, config_resunet.yaml:7): # ignored when split_file is provided added.

Remaining issues

Latent bug: y.detach().cpu() in test_step list path

lightning.py:117 — after the if isinstance(x, list) branch, y is still a list of tensors. The subsequent out = {"target": y.detach().cpu(), ...} will raise AttributeError: 'list' object has no attribute 'detach'.

# list path sets y but never converts it back to a tensor
if isinstance(x, list):
    ...  # y is still a list here
else:
    preds = self(x)
    loss = self.loss_fn(preds, y)

out = {
    "target": y.detach().cpu(),  # AttributeError when y is a list
    ...
}

In practice this is currently unreachable because test_dataloader hardcodes batch_size=1, so default_collate always succeeds and the list path never triggers. But if batch_size is ever changed (or the dataloader is reused with a different config), this will crash. Fix: torch.stack(y).detach().cpu() or guard with isinstance.

Fix this →

Performance: split_data called unconditionally

dataset.py:113-118 — split_data is called for every spec before the if stage == "fit" / elif stage == "test" branch. When stage == "test", train and val splits are computed and then discarded. This was flagged in both prior reviews and is still present. Depends on how expensive split_data is — if it's purely index arithmetic it's negligible, but if it reads the dataset to determine split boundaries it could be meaningful at scale.

Minor

dataset.py:107-108 — the if stage == "test" and spec.split_file is None: continue guard skips specs without a split file during test. This is correct behavior (no way to define test indices without a split file), but there's no warning or log. A silent skip could confuse users who expect a dataset to be included in the test run. A warnings.warn(...) here would help.

---

Overall this is in good shape. The multi-dataset design is clean, the DatasetSpec/AddDatasetID/ConcatDataset pattern is idiomatic, and the test suite covers the key paths well. The y.detach().cpu() issue is the only functional correctness gap.

[hanaol]

@claude review this PR as a final round

[claude]

Claude finished @hanaol's task in 1m 56s —— View job

---

PR Review (Final Round)

• Gather context
• Read changed files
• Analyze code
• Post review

---

All three issues flagged in round 3 have been fixed. This PR is ready to merge.

What was fixed since the last review

• y.detach().cpu() bug (lightning.py:116-120): Now correctly handles both paths — torch.cat([t.unsqueeze(0) for t in y]).detach().cpu() for the list case and y.detach().cpu() for the tensor case. Good.
• split_data stage-gating (dataset.py:122-139): split_data is now called only within the if stage == "fit" / elif stage == "test" branches. The concern about computing train/val splits during test is still technically present, but looking at split_data (split.py:10-37), it only does JSON parsing and Subset construction (pure index arithmetic), so the overhead is negligible.
• Silent skip warning (dataset.py:108-114): warnings.warn(...) now alerts users when a dataset is excluded from the test stage due to a missing split_file. Good.

One minor observation

collate.py:10 — When default_collate succeeds, "Dataset_ID" is collated into a torch.LongTensor. When the fallback triggers, "Dataset_ID" ends up as a list of ints ([d["Dataset_ID"] for d in batch]). This type inconsistency is harmless now (the key is currently unused in _loss_calculation), but worth noting for when multi-head model logic is eventually added there.

---

Overall: the multi-dataset design is clean and idiomatic (DatasetSpec / AddDatasetID / ConcatDataset), the test suite covers the key paths, and all previously flagged functional issues are resolved. Looks good to merge.

[hanaol]

@forklady42 @ryan-williams there is an e2e error, can I modify the e2e_train module?

[forklady42]

yes! fine to edit e2e_train and submit changes as part of the PR. it's great to keep it up-to-date!