Device-resident batch sampling for INR tomography reconstruction

src/quantem/tomography/dataset_models.py

@@ -440,6 +440,96 @@ def to(self, device: str | torch.device):
         self.device = device
 
 
+class DeviceBatchSampler:
+    """Epoch iterator that builds INR training batches directly on a device.
+
+    Replaces the per-pixel DataLoader path: the tilt stack and angles are
+    made resident on ``device`` once, so producing a batch is index
+    arithmetic plus two tensor lookups instead of ``batch_size`` Python
+    ``__getitem__`` calls, a collate, and a host-to-device copy per step.
+    On a GPU this removes the CPU dataloader bottleneck entirely.
+
+    Yields the same batch dicts as ``TomographyINRDataset.__getitem__``
+    under a DataLoader collate (``projection_idx``, ``pixel_i``,
+    ``pixel_j``, ``phi``, ``target_value``), with the train loader's
+    ``drop_last=True`` semantics.
+
+    Distributed runs: pass ``rank``/``world_size`` and every rank derives
+    the *same* epoch permutation from ``seed + epoch`` (CPU generator, so
+    it is identical across ranks and reproducible), then takes an
+    equal-size contiguous shard — equal so per-rank batch counts match and
+    DDP gradient sync cannot hang on a ragged tail. The training loop's
+    ``sampler.set_epoch(epoch)`` drives reshuffling, exactly like
+    ``DistributedSampler``; without ``set_epoch`` the epoch advances
+    automatically on each ``__iter__``.
+    """
+
+    def __init__(
+        self,
+        dset: "TomographyINRDataset",
+        batch_size: int,
+        device: torch.device | str,
+        indices: torch.Tensor | None = None,
+        shuffle: bool = True,
+        rank: int = 0,
+        world_size: int = 1,
+        seed: int = 0,
+    ):
+        self.batch_size = batch_size
+        self.device = torch.device(device)
+        self.shuffle = shuffle
+        self.rank = rank
+        self.world_size = world_size
+        self.seed = seed
+        self._epoch = 0
+        self._stack = dset.tilt_stack.to(self.device)
+        self._angles = dset.tilt_angles.to(self.device)
+        # __getitem__ decodes flat indices with shape[1] for both rows and
+        # columns; replicate it exactly.
+        self._s1 = dset.tilt_stack.shape[1]
+        self._s2 = dset.tilt_stack.shape[2]
+        if indices is None:
+            indices = torch.arange(len(dset), dtype=torch.int64)
+        self._indices = indices.to(self.device)
+        self._per_rank = len(self._indices) // world_size
+
+    def set_epoch(self, epoch: int) -> None:
+        """Set the epoch used to seed this epoch's shared permutation."""
+        self._epoch = epoch
+
+    def __len__(self) -> int:
+        return self._per_rank // self.batch_size  # drop_last=True
+
+    def _epoch_shard(self) -> torch.Tensor:
+        idx = self._indices
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.seed + self._epoch)
+            perm = torch.randperm(len(idx), generator=g).to(self.device)
+            idx = idx[perm]
+            self._epoch += 1  # auto-advance; set_epoch overrides per epoch
+        if self.world_size > 1:
+            idx = idx[self.rank * self._per_rank : (self.rank + 1) * self._per_rank]
+        return idx
+
+    def __iter__(self):
+        idx = self._epoch_shard()
+        per_proj = self._s1 * self._s2
+        for k in range(len(self)):
+            sel = idx[k * self.batch_size : (k + 1) * self.batch_size]
+            proj = sel // per_proj
+            rem = sel - proj * per_proj
+            pixel_i = rem // self._s1
+            pixel_j = rem - pixel_i * self._s1
+            yield {
+                "projection_idx": proj,
+                "pixel_i": pixel_i,
+                "pixel_j": pixel_j,
+                "phi": self._angles[proj],
+                "target_value": self._stack[proj, pixel_i, pixel_j],
+            }
+
+
 class TomographyINRDataset(TomographyDatasetConstraints, Dataset):
     """
     Dataset class for INR-based tomography.

src/quantem/tomography/tomography.py

@@ -17,6 +17,7 @@
     DatasetConstraintParams,
     DatasetConstraintsType,
     DatasetModelType,
+    DeviceBatchSampler,
     TomographyINRDataset,
     TomographyPixDataset,
 )
@@ -145,14 +146,7 @@ def reconstruct(
                     self.scheduler_params = scheduler_params
                     self.set_schedulers(self.scheduler_params, num_iter=num_iter)
 
-            self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = (
-                self.setup_dataloader(
-                    self.dset,
-                    batch_size,
-                    num_workers=num_workers,
-                    val_fraction=val_fraction,
-                )
-            )
+            self._setup_recon_dataloaders(batch_size, num_workers, val_fraction)
 
         # Type check for INR-based reconstruction
         if not isinstance(self.dset, TomographyINRDataset):
@@ -285,10 +279,14 @@ def reconstruct(
                     val_loss = torch.tensor(0.0, device=self.device)
 
                     for batch in self.val_dataloader:
+                        # Match the training pass (enabled=False): bf16 autocast
+                        # breaks the so3 pose solve (lu_factor has no BFloat16
+                        # kernel) and would make the val loss inconsistent with
+                        # the fp32 training loss it is compared to.
                         with torch.autocast(
                             device_type=self.device.type,
                             dtype=torch.bfloat16,
-                            enabled=True,
+                            enabled=False,
                         ):
                             all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
 
@@ -411,6 +409,49 @@ def _rebuild_dataloader(self, batch_size: int, num_workers: int, val_fraction: f
         """
         Rebuilds the dataloader due to persistent workers error when reloading the object.
         """
+        self._setup_recon_dataloaders(batch_size, num_workers, val_fraction)
+
+    def _setup_recon_dataloaders(self, batch_size: int, num_workers: int, val_fraction: float):
+        """Build the train/val batch iterators.
+
+        INR datasets use ``DeviceBatchSampler`` — batches are built with
+        tensor ops on the compute device from a device-resident tilt stack,
+        removing the per-pixel ``__getitem__`` / collate / H2D-copy
+        dataloader bottleneck (``num_workers`` is ignored on this path).
+        Distributed runs shard the same seeded epoch permutation across
+        ranks (DistributedSampler semantics; the loop's ``set_epoch`` drives
+        reshuffling). Non-INR datasets keep the DataLoader path.
+        """
+        if isinstance(self.dset, TomographyINRDataset):
+            n = len(self.dset)
+            n_val = int(n * val_fraction)
+            # Fixed-seed split: identical across DDP ranks (no train/val
+            # leakage between ranks) and stable across save/reload, so a
+            # resumed run keeps validating on the same held-out pixels.
+            split_gen = torch.Generator()
+            split_gen.manual_seed(0)
+            perm = torch.randperm(n, generator=split_gen)
+            ddp = dict(rank=self.global_rank, world_size=self.world_size)
+            self.dataloader = DeviceBatchSampler(
+                self.dset, batch_size, self.device, indices=perm[n_val:], **ddp
+            )
+            # The val sampler keeps its own device-resident copy of the tilt
+            # stack; acceptable, since val_fraction > 0 is the rare case.
+            val = (
+                DeviceBatchSampler(
+                    self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False, **ddp
+                )
+                if n_val > 0
+                else None
+            )
+            # A per-rank val shard smaller than one batch would divide by
+            # zero in the val-loss average.
+            self.val_dataloader = val if val is not None and len(val) > 0 else None
+            # The training loop calls set_epoch on self.sampler.
+            self.sampler = self.dataloader
+            self.val_sampler = None
+            return
+
         self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = (
             self.setup_dataloader(
                 self.dset,

tests/tomography/test_device_batch_sampler.py

@@ -0,0 +1,142 @@
+"""Tests for ``DeviceBatchSampler`` and its wiring in ``Tomography.reconstruct``.
+
+The sampler must yield batches identical in content to the per-pixel
+DataLoader path (same keys, same index decode as
+``TomographyINRDataset.__getitem__``), cover each pixel exactly once per
+epoch (minus the dropped tail batch), and respect the train/val index
+split. Device-independent, so everything here runs on CPU.
+"""
+
+import numpy as np
+import torch
+
+from quantem.tomography.dataset_models import DeviceBatchSampler, TomographyINRDataset
+
+
+def _dset(n_proj=4, n=10, seed=0):
+    rng = np.random.default_rng(seed)
+    stack = rng.random((n_proj, n, n)).astype(np.float32)
+    angles = np.linspace(-60, 60, n_proj).astype(np.float32)
+    return TomographyINRDataset.from_data(tilt_stack=stack, tilt_angles=angles)
+
+
+def test_batches_match_getitem():
+    dset = _dset()
+    sampler = DeviceBatchSampler(dset, batch_size=37, device="cpu", shuffle=False)
+    seen = 0
+    for batch in sampler:
+        for k in range(len(batch["target_value"])):
+            item = dset[seen + k]
+            for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value"):
+                torch.testing.assert_close(
+                    batch[key][k], item[key].to(batch[key].dtype), rtol=0, atol=0
+                )
+        seen += len(batch["target_value"])
+
+
+def test_epoch_covers_indices_once_with_drop_last():
+    dset = _dset()
+    n = len(dset)
+    batch_size = 64
+    sampler = DeviceBatchSampler(dset, batch_size=batch_size, device="cpu", shuffle=True)
+    assert len(sampler) == n // batch_size
+    per_proj = dset.tilt_stack.shape[1] * dset.tilt_stack.shape[2]
+    flat = []
+    for batch in sampler:
+        assert len(batch["target_value"]) == batch_size
+        flat.append(
+            batch["projection_idx"] * per_proj
+            + batch["pixel_i"] * dset.tilt_stack.shape[1]
+            + batch["pixel_j"]
+        )
+    flat = torch.cat(flat)
+    assert flat.unique().numel() == flat.numel()  # no repeats within an epoch
+    assert flat.numel() == len(sampler) * batch_size
+
+
+def test_shuffle_changes_order_between_epochs():
+    dset = _dset()
+    sampler = DeviceBatchSampler(dset, batch_size=50, device="cpu", shuffle=True)
+    first = next(iter(sampler))["target_value"]
+    second = next(iter(sampler))["target_value"]
+    assert not torch.equal(first, second)
+
+
+def test_val_split_is_disjoint():
+    dset = _dset()
+    n = len(dset)
+    perm = torch.randperm(n)
+    n_val = n // 10
+    train = DeviceBatchSampler(dset, 32, "cpu", indices=perm[n_val:])
+    val = DeviceBatchSampler(dset, 32, "cpu", indices=perm[:n_val], shuffle=False)
+    assert len(train._indices) + len(val._indices) == n
+    assert torch.cat([train._indices, val._indices]).unique().numel() == n
+
+
+def test_reconstruct_uses_sampler_single_process():
+    """Smoke: the single-process reconstruct path builds DeviceBatchSamplers.
+
+    Object models that go through ``setup_distributed`` must be built on CUDA
+    when a CUDA device exists, so pick the device accordingly.
+    """
+    from quantem.core.ml.models.kplanes import KPlanesTILTED
+    from quantem.core.ml.optimizer_mixin import OptimizerParams
+    from quantem.tomography.object_models import ObjectINR
+    from quantem.tomography.tomography import Tomography
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    dset = _dset(n_proj=3, n=8)
+    model = KPlanesTILTED(M_features=2, resolution=(8, 8, 8), multiscale_res_multipliers=[1], T=1)
+    obj = ObjectINR.from_model(model, shape=(8, 8, 8), device=device)
+    tomo = Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False)
+    tomo.reconstruct(
+        num_iter=1,
+        batch_size=32,
+        num_workers=0,
+        val_fraction=0.25,
+        optimizer_params={
+            "object": {"default": OptimizerParams.Adam(lr=1e-3)},
+            "pose": OptimizerParams.Adam(lr=1e-3),
+        },
+    )
+    assert isinstance(tomo.dataloader, DeviceBatchSampler)
+    assert isinstance(tomo.val_dataloader, DeviceBatchSampler)
+    # the loop drives epoch reshuffling through set_epoch on self.sampler
+    assert tomo.sampler is tomo.dataloader
+
+
+def _flat(batch, s1, s2):
+    return batch["projection_idx"] * (s1 * s2) + batch["pixel_i"] * s1 + batch["pixel_j"]
+
+
+def test_ddp_shards_are_disjoint_and_equal():
+    dset = _dset()
+    s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2]
+    world = 4
+    shards = []
+    for rank in range(world):
+        sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=world)
+        sampler.set_epoch(3)
+        flats = torch.cat([_flat(b, s1, s2) for b in sampler])
+        shards.append(flats)
+        assert len(sampler) == (len(dset) // world) // 32  # equal on every rank
+    allv = torch.cat(shards)
+    assert allv.unique().numel() == allv.numel()  # no pixel on two ranks
+
+
+def test_ddp_epoch_permutation_shared_and_reproducible():
+    dset = _dset()
+    s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2]
+
+    def epoch_flats(rank, epoch):
+        sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=2)
+        sampler.set_epoch(epoch)
+        return torch.cat([_flat(b, s1, s2) for b in sampler])
+
+    # same (rank, epoch) on a fresh instance -> identical batches
+    torch.testing.assert_close(epoch_flats(0, 7), epoch_flats(0, 7), rtol=0, atol=0)
+    # different epoch -> different order
+    assert not torch.equal(epoch_flats(0, 7), epoch_flats(0, 8))
+    # ranks of the same epoch are disjoint
+    both = torch.cat([epoch_flats(0, 7), epoch_flats(1, 7)])
+    assert both.unique().numel() == both.numel()