depth_fm.data.datamodule

datamodule

Fixed LitData DataModule — works around StreamingDataLoader re-iteration deadlock.

THE BUG

LitData's StreamingDataLoader doesn't properly reset its internal iterator state when re-iterated (GitHub Issues #316, #213, #452). In DDP, this causes ranks to get different batch counts on the 2nd+ validation, deadlocking on the next sync_dist or NCCL collective.

THE FIX (two options, both provided):

Option A (default): Use torch.utils.data.DataLoader for val/test. StreamingDataset is an IterableDataset and works fine with the regular DataLoader. You lose StreamingDataLoader's prefetch optimizations, but val/test are small — this costs ~1 second per validation.

Option B: Recreate StreamingDataset + StreamingDataLoader from scratch each time val_dataloader() is called. This avoids the stale-state bug by never re-iterating the same object. Slightly more overhead from worker startup.

Drop-in replacement for litdata_datamodule_1.py — same API, same config.

MarsStreamingDataset

MarsStreamingDataset(input_dir: str, is_train: bool = False, random_flip: bool = True, brightness_jitter: float = 0.1, shuffle: bool = False, drop_last: bool = True, seed: int = 42)

Bases: StreamingDataset

StreamingDataset subclass with on-the-fly augmentations.

Source code in src/depth_fm/data/datamodule.py

def __init__(
        self,
        input_dir: str,
        is_train: bool = False,
        random_flip: bool = True,
        brightness_jitter: float = 0.1,
        shuffle: bool = False,
        drop_last: bool = True,
        seed: int = 42,
):
    super().__init__(
        input_dir=input_dir,
        shuffle=shuffle,
        drop_last=drop_last,
        seed=seed,
    )
    self.is_train = is_train
    self.random_flip = random_flip and is_train
    self.brightness_jitter = brightness_jitter if is_train else 0.0

MarsDepthFMDataModule

MarsDepthFMDataModule(config: DictConfig)

Bases: LightningDataModule

Lightning DataModule with LitData re-iteration deadlock fix.

Key change: val/test use torch.utils.data.DataLoader instead of StreamingDataLoader. The StreamingDataset works with both — we only lose LitData's prefetch optimizations, which are irrelevant for the small val/test sets (~25 batches).

Train still uses StreamingDataLoader for its shuffle + prefetch benefits.

Source code in src/depth_fm/data/datamodule.py

def __init__(self, config: DictConfig):
    super().__init__()
    self.config = config
    self.hc = config.data.hirise
    self.tc = config.training

    self.cache_hash = litdata_cache_key(config)
    self.litdata_root = litdata_cache_root(config)

    self._train_dataset = None
    # val/test datasets stored only if strategy == "torch"
    self._val_dataset = None
    self._test_dataset = None