From 4d4f32577f6cf92e82194f5010a7c7793c5765a1 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Tue, 9 Jun 2026 23:00:20 -0700 Subject: [PATCH 01/19] Build INR training batches on the compute device for single-process runs The per-pixel DataLoader path costs ~43 ms/batch on CPU (8192 Python __getitem__ calls + collate + H2D copy) while the GPU step itself takes ~20 ms, so single-GPU reconstructions idle the GPU half the time. DeviceBatchSampler keeps the tilt stack resident on the device and builds each batch with index arithmetic and two tensor lookups, yielding the same batch dicts (and drop_last/val-split semantics) as the DataLoader path. DDP runs keep the DataLoader + DistributedSampler path. Also disable bf16 autocast in the validation loop to match the training pass: the so3 pose solve (lu_factor) has no BFloat16 kernel, and a bf16 val loss is not comparable to the fp32 train loss. --- src/quantem/tomography/dataset_models.py | 60 ++++++++++ src/quantem/tomography/tomography.py | 48 ++++++-- tests/tomography/test_device_batch_sampler.py | 104 ++++++++++++++++++ 3 files changed, 203 insertions(+), 9 deletions(-) create mode 100644 tests/tomography/test_device_batch_sampler.py diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 2255636c..00419695 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -440,6 +440,66 @@ 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 for single-process runs: 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. + """ + + def __init__( + self, + dset: "TomographyINRDataset", + batch_size: int, + device: torch.device | str, + indices: torch.Tensor | None = None, + shuffle: bool = True, + ): + self.batch_size = batch_size + self.device = torch.device(device) + self.shuffle = shuffle + 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) + + def __len__(self) -> int: + return len(self._indices) // self.batch_size # drop_last=True + + def __iter__(self): + idx = self._indices + if self.shuffle: + idx = idx[torch.randperm(len(idx), device=self.device)] + 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. diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index f86095ad..2fdbb7cf 100644 --- a/src/quantem/tomography/tomography.py +++ b/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,38 @@ 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. + + Single-process runs 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). + Multi-process (DDP) runs keep the DataLoader + DistributedSampler + path. + """ + if self.world_size == 1 and isinstance(self.dset, TomographyINRDataset): + n = len(self.dset) + n_val = int(n * val_fraction) + perm = torch.randperm(n) + self.dataloader = DeviceBatchSampler( + self.dset, batch_size, self.device, indices=perm[n_val:] + ) + # The val sampler keeps its own device-resident copy of the tilt + # stack; acceptable, since val_fraction > 0 is the rare case. + self.val_dataloader = ( + DeviceBatchSampler( + self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False + ) + if n_val > 0 + else None + ) + self.sampler = None + self.val_sampler = None + return + self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = ( self.setup_dataloader( self.dset, diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py new file mode 100644 index 00000000..e40b1284 --- /dev/null +++ b/tests/tomography/test_device_batch_sampler.py @@ -0,0 +1,104 @@ +"""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) + assert tomo.sampler is None From 9e7e505a30f45850d6117b231ac580e444e3bef0 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Tue, 9 Jun 2026 23:30:59 -0700 Subject: [PATCH 02/19] Shard DeviceBatchSampler across DDP ranks Every rank derives the same epoch permutation from seed + epoch (CPU generator, identical across ranks and reproducible) and takes an equal-size contiguous shard, with the ragged tail dropped so per-rank batch counts always match and gradient sync cannot hang. The training loop's existing sampler.set_epoch call drives reshuffling, exactly like DistributedSampler; single-process runs auto-advance the epoch instead. The train/val split now uses a fixed-seed generator: identical across ranks (no leakage between a rank's train shard and another's val shard) and stable across save/reload, so resumed runs keep validating on the same held-out pixels. Verified with a 2-GPU torchrun run: equal batch counts on both ranks, finite identical reduced losses, no deadlock. --- src/quantem/tomography/dataset_models.py | 48 +++++++++++++++---- src/quantem/tomography/tomography.py | 33 ++++++++----- tests/tomography/test_device_batch_sampler.py | 40 +++++++++++++++- 3 files changed, 100 insertions(+), 21 deletions(-) diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 00419695..af69d205 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -443,17 +443,25 @@ def to(self, device: str | torch.device): class DeviceBatchSampler: """Epoch iterator that builds INR training batches directly on a device. - Replaces the per-pixel DataLoader path for single-process runs: 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. + 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__( @@ -463,10 +471,17 @@ def __init__( 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 @@ -476,14 +491,29 @@ def __init__( 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 len(self._indices) // self.batch_size # drop_last=True + return self._per_rank // self.batch_size # drop_last=True - def __iter__(self): + def _epoch_shard(self) -> torch.Tensor: idx = self._indices if self.shuffle: - idx = idx[torch.randperm(len(idx), device=self.device)] + 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] diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index 2fdbb7cf..4a68edc5 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -414,30 +414,41 @@ def _rebuild_dataloader(self, batch_size: int, num_workers: int, val_fraction: f def _setup_recon_dataloaders(self, batch_size: int, num_workers: int, val_fraction: float): """Build the train/val batch iterators. - Single-process runs 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 + 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). - Multi-process (DDP) runs keep the DataLoader + DistributedSampler - 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 self.world_size == 1 and isinstance(self.dset, TomographyINRDataset): + if isinstance(self.dset, TomographyINRDataset): n = len(self.dset) n_val = int(n * val_fraction) - perm = torch.randperm(n) + # 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:] + 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. - self.val_dataloader = ( + val = ( DeviceBatchSampler( - self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False + self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False, **ddp ) if n_val > 0 else None ) - self.sampler = 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 diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py index e40b1284..47bce96f 100644 --- a/tests/tomography/test_device_batch_sampler.py +++ b/tests/tomography/test_device_batch_sampler.py @@ -101,4 +101,42 @@ def test_reconstruct_uses_sampler_single_process(): ) assert isinstance(tomo.dataloader, DeviceBatchSampler) assert isinstance(tomo.val_dataloader, DeviceBatchSampler) - assert tomo.sampler is None + # 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() From 5da8a7e734326d58ed14c13ca6d2a5cf7dde840b Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 01:48:48 -0700 Subject: [PATCH 03/19] Fix INR dataset indexing for non-square tilt images and TV soft-constraint crash - TomographyINRDataset.__getitem__ decomposed pixel indices by shape[1] (height) instead of shape[2] (width), and __len__ used max(shape)^2 instead of H*W; both corrupted the pixel mapping and overcounted the dataset for rectangular tilt images. - ObjectPixelated.apply_soft_constraints built soft_loss as a leaf tensor with requires_grad=True and then added the TV loss in-place, which raises RuntimeError whenever tv_vol > 0. Build it without requires_grad and add out-of-place so the loss backprops through get_tv_loss. --- src/quantem/tomography/dataset_models.py | 7 +++---- src/quantem/tomography/object_models.py | 6 ++---- tests/tomography/test_dataset_models.py | 19 +++++++++++++++++++ tests/tomography/test_object_models.py | 11 +++++++++++ 4 files changed, 35 insertions(+), 8 deletions(-) diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 2255636c..515314f4 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -611,8 +611,8 @@ def __getitem__( projection_idx = actual_idx // (self.tilt_stack.shape[1] * self.tilt_stack.shape[2]) remaining = actual_idx % (self.tilt_stack.shape[1] * self.tilt_stack.shape[2]) - pixel_i = remaining // self.tilt_stack.shape[1] - pixel_j = remaining % self.tilt_stack.shape[1] + pixel_i = remaining // self.tilt_stack.shape[2] + pixel_j = remaining % self.tilt_stack.shape[2] return { "projection_idx": torch.tensor(projection_idx), @@ -628,8 +628,7 @@ def __len__( """ Returns the number of pixels in the tilt stack. """ - N = max(self.tilt_stack.shape) - return self.tilt_stack.shape[0] * N * N + return self.tilt_stack.shape[0] * self.tilt_stack.shape[1] * self.tilt_stack.shape[2] def to(self, device: torch.device | str): self._z1_params = nn.Parameter(self._z1_angles.to(device)) diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index 817e7ed5..ccdfd539 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -421,12 +421,10 @@ def apply_hard_constraints( def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor: assert ctx.obj is not None, "ObjectPixelated requires ctx.obj to be set" - soft_loss = torch.tensor( - 0.0, device=ctx.obj.device, dtype=ctx.obj.dtype, requires_grad=True - ) + soft_loss = torch.tensor(0.0, device=ctx.obj.device, dtype=ctx.obj.dtype) if self.constraints.tv_vol > 0: tv_loss = self.get_tv_loss(ctx) - soft_loss += tv_loss + soft_loss = soft_loss + tv_loss return soft_loss # --- Forward method --- diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py index 04abd3b5..5534ec36 100644 --- a/tests/tomography/test_dataset_models.py +++ b/tests/tomography/test_dataset_models.py @@ -180,6 +180,25 @@ def test_getitem_index_mapping(self): assert int(item["pixel_j"]) == 1 assert torch.isclose(item["target_value"], d.tilt_stack[1, 1, 1]) + def test_len_and_getitem_non_square(self): + # Regression: pixel_i/pixel_j must decompose by the width (shape[2]) and __len__ + # by H*W, not max(shape)^2 -- both were wrong for rectangular tilt images. + rng = np.random.default_rng(0) + stack = rng.random((3, 4, 6)).astype(np.float32) # nang=3, H=4, W=6 + d = TomographyINRDataset.from_data(stack, np.linspace(-60, 60, 3, dtype="f4")) + assert len(d) == 3 * 4 * 6 + # idx = proj*(H*W) + i*W + j = 1*24 + 2*6 + 3 = 39 + item = d[39] + assert int(item["projection_idx"]) == 1 + assert int(item["pixel_i"]) == 2 + assert int(item["pixel_j"]) == 3 + assert torch.isclose(item["target_value"], d.tilt_stack[1, 2, 3]) + # every index in range maps to a valid pixel + last = d[len(d) - 1] + assert int(last["projection_idx"]) == 2 + assert int(last["pixel_i"]) == 3 + assert int(last["pixel_j"]) == 5 + def test_save_load_parameters_roundtrip(self, tmp_path): angles = np.linspace(-60, 60, 5, dtype="f4") d = TomographyINRDataset.from_data(_stack(), angles) diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py index c4216827..c59f4d28 100644 --- a/tests/tomography/test_object_models.py +++ b/tests/tomography/test_object_models.py @@ -148,6 +148,17 @@ def test_soft_constraint_zero_when_tv_off(self): obj.constraints.tv_vol = 0.0 assert float(obj.apply_soft_constraints(ctx).detach()) == 0.0 + def test_soft_constraint_with_tv_on_backprops(self): + # Regression: soft_loss was a leaf tensor with requires_grad=True, so the + # in-place `+= tv_loss` raised RuntimeError whenever tv_vol > 0. + ctx = ReconstructionContext(obj=torch.rand(8, 8, 8).requires_grad_(True)) + obj = ObjectPixelated.from_uniform(shape=(8, 8, 8), device="cpu") + obj.constraints.tv_vol = 1.0 + loss = obj.apply_soft_constraints(ctx) + assert torch.isfinite(loss) and loss > 0.0 + loss.backward() + assert ctx.obj.grad is not None + class TestFactoryGuard: def test_objectbase_requires_token(self): From 48b4213d936057fa54981d0a7120ed0a867ce6e4 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 01:52:27 -0700 Subject: [PATCH 04/19] Speed up KPlanesTILTED feature interpolation and volume TV loss - interpolate_ms_features_tilted: build the per-plane coordinate pairs with unbind/stack views instead of allocating an index tensor on device every call (a host-to-device copy) and running an (T, 3, B, 3) expand+gather. Bitwise-identical output; ~1.1x on the isolated forward (B=4096, T=8, 3 scales) plus one fewer H2D transfer per model call. - ObjectTensorDecomp.get_volume_tv_loss: evaluate the base points and the three axis-shifted copies in one batched forward (4N points) instead of four sequential model calls. Bitwise-identical loss; 3.1x faster forward and 3.3x with backward (10k samples, KPlanesTILTED backbone). --- src/quantem/core/ml/models/kplanes.py | 21 ++++++++------ src/quantem/tomography/object_models.py | 38 ++++++++++++------------- 2 files changed, 32 insertions(+), 27 deletions(-) diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py index cdf55261..67ea3089 100644 --- a/src/quantem/core/ml/models/kplanes.py +++ b/src/quantem/core/ml/models/kplanes.py @@ -169,6 +169,7 @@ class KPlanes(PPLR, TensorDecompositionModel): """ K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241 """ + def __init__( self, # Grid parameters @@ -320,14 +321,18 @@ def interpolate_ms_features_tilted( rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) # Build (T, 3, B, 2) coords for planes XY, ZX, YZ in one shot. - # index_select is faster and cleaner than advanced indexing with python lists. - # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2) - idx = torch.tensor([[0, 1], [2, 0], [1, 2]], device=pts.device) # (3, 2) - # rotated: (T, B, 3) -> gather along last dim with idx (3, 2) - # Result: (T, 3, B, 2) - coords = ( - rotated.unsqueeze(1).expand(T, 3, B, 3).gather(-1, idx.view(1, 3, 1, 2).expand(T, 3, B, 2)) - ) + # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2). Stacking views avoids the + # per-call index-tensor allocation (a host-to-device copy) and the + # (T, 3, B, 3) expand+gather this used to do. + x, y, z = rotated.unbind(-1) # each (T, B) + coords = torch.stack( + ( + torch.stack((x, y), dim=-1), + torch.stack((z, x), dim=-1), + torch.stack((y, z), dim=-1), + ), + dim=1, + ) # (T, 3, B, 2) # Flatten (T, 3) -> 3*T so it matches grid's first dim, and add the H_out=1 axis coord_tensor = coords.reshape(3 * T, B, 1, 2) # (3T, B, 1, 2) diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index ccdfd539..aceb562e 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -983,25 +983,25 @@ def get_volume_tv_loss(self, coords: torch.Tensor) -> torch.Tensor: model = _unwrap(self.model) h = 2.0 / min(model.resolution) - pred = model(tv_coords) - if isinstance(pred, tuple): - pred = pred[0] - if pred.dim() == 1: - pred = pred.unsqueeze(-1) # (N, 1) - - grads = [] - for axis in range(3): - offset = torch.zeros(3, device=tv_coords.device) - offset[axis] = h - shifted_pred = self.model(tv_coords + offset) - if isinstance(shifted_pred, tuple): - shifted_pred = shifted_pred[0] - if shifted_pred.dim() == 1: - shifted_pred = shifted_pred.unsqueeze(-1) - grads.append((shifted_pred - pred) / h) # (N, 1) - - grad_stack = torch.stack(grads, dim=-1) # (N, C, 3) - grad_norm = torch.norm(grad_stack, dim=-1) # (N, C) + # Evaluate the base points and the three axis-shifted copies in a single + # batched forward (4N points) instead of 4 sequential model calls. + offsets = h * torch.eye(3, device=tv_coords.device, dtype=tv_coords.dtype) # (3, 3) + all_coords = torch.cat( + [tv_coords, (tv_coords.unsqueeze(0) + offsets.unsqueeze(1)).reshape(-1, 3)] + ) # (4N, 3) + + all_pred = model(all_coords) + if isinstance(all_pred, tuple): + all_pred = all_pred[0] + if all_pred.dim() == 1: + all_pred = all_pred.unsqueeze(-1) # (4N, 1) + + n = tv_coords.shape[0] + pred = all_pred[:n] # (N, C) + shifted_pred = all_pred[n:].view(3, n, -1) # (3, N, C) + + grad_stack = (shifted_pred - pred.unsqueeze(0)) / h # (3, N, C) + grad_norm = torch.norm(grad_stack, dim=0) # (N, C) return self.constraints.tv_vol * grad_norm.mean() From a16ea3862c8eac0004f9776b04371bacadc15359 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 02:25:08 -0700 Subject: [PATCH 05/19] Fix ObjectINR out-of-volume masking along z and soft-constraint device crash - ObjectINR.forward masked densities outside [-1, 1] in x and y only, so tilted rays whose sample points leave the volume along z still picked up extrapolated density and biased the integrated projections at high tilt. The mask now bounds all three axes. - ObjectINR.apply_soft_constraints read ctx.coords.device before the coords-is-None assert, raising AttributeError instead of returning a zero loss when no soft constraints are active and no coords are passed. Fall back to the model device when coords are absent. --- src/quantem/tomography/object_models.py | 10 ++++++++-- tests/tomography/test_object_models.py | 21 ++++++++++++++++++--- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index 817e7ed5..504cc3a7 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -542,7 +542,8 @@ def apply_soft_constraints( self, ctx: ReconstructionContext, ) -> torch.Tensor: - soft_loss = torch.tensor(0.0, device=ctx.coords.device) + device = ctx.coords.device if ctx.coords is not None else self._device + soft_loss = torch.tensor(0.0, device=device) if self.constraints.tv_vol > 0: assert ctx.coords is not None, ( "coords must be provided for INR object model to compute the TV loss" @@ -662,7 +663,12 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor: if all_densities.dim() > 1: all_densities = all_densities.squeeze(-1) valid_mask = ( - (coords[:, 0] >= -1) & (coords[:, 0] <= 1) & (coords[:, 1] >= -1) & (coords[:, 1] <= 1) + (coords[:, 0] >= -1) + & (coords[:, 0] <= 1) + & (coords[:, 1] >= -1) + & (coords[:, 1] <= 1) + & (coords[:, 2] >= -1) + & (coords[:, 2] <= 1) ).float() if all_densities.dim() > 1: diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py index c4216827..9123300e 100644 --- a/tests/tomography/test_object_models.py +++ b/tests/tomography/test_object_models.py @@ -205,10 +205,25 @@ def _obj(self, device, n=16): def test_forward_masks_out_of_range(self, torch_device): obj = self._obj(torch_device) - coords = torch.tensor([[0.0, 0.0, 0.0], [5.0, 0.0, 0.0]], device=torch_device) + # Regression: the mask used to check x and y only, so tilted rays leaving the + # volume along z still contributed (extrapolated) density to the integral. + coords = torch.tensor( + [[0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 5.0]], + device=torch_device, + ) out = obj.forward(coords) - assert out.shape[0] == 2 - assert float(out[1].detach()) == 0.0 # x=5 is outside [-1, 1] -> masked to zero + assert out.shape[0] == 4 + for i in (1, 2, 3): # each axis outside [-1, 1] -> masked to zero + assert float(out[i].detach()) == 0.0 + + def test_soft_constraints_without_coords(self, torch_device): + # Regression: soft_loss was created on ctx.coords.device before the None + # check, raising AttributeError when no constraints are active. + obj = self._obj(torch_device) + obj.constraints.tv_vol = 0.0 + obj.constraints.sparsity = 0.0 + loss = obj.apply_soft_constraints(ReconstructionContext()) + assert float(loss.detach()) == 0.0 def test_apply_hard_constraints_positivity(self, torch_device): obj = self._obj(torch_device) From 0eb870f4ad640589c9e313a6bdf97c6a01049323 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 02:28:43 -0700 Subject: [PATCH 06/19] Reduce INR dataloader and ray-transform overhead - TomographyINRDataset.__getitem__ wrapped each of the three index fields in torch.tensor(), allocating three scalar tensors per item on the dataloader hot path; default_collate builds the same int64 batch tensors from plain ints. 2.9x faster batch loading (12.1 -> 4.2 ms/batch, batch_size=1024, 60x256x256 stack), which dominated the per-batch GPU compute (~0.4 ms). - transform_batch_rays applied the three Euler rotations as nine elementwise passes over the full (B, S) ray tensors. Compose them into one (B, 3, 3) matrix and apply with a single batched matmul: same result to 4e-7 (float32 op reordering), up to 1.15x at large batches and far fewer kernel launches. --- src/quantem/tomography/dataset_models.py | 69 ++++++++++++------------ 1 file changed, 36 insertions(+), 33 deletions(-) diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 515314f4..9ccce8c1 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -550,37 +550,37 @@ def transform_batch_rays( shift_x_norm = (shifts[:, 0:1] * sampling_rate * 2) / (N - 1) shift_y_norm = (shifts[:, 1:2] * sampling_rate * 2) / (N - 1) - rays_x = rays[:, :, 0] - shift_x_norm - rays_y = rays[:, :, 1] - shift_y_norm - rays_z = rays[:, :, 2] - - theta = torch.deg2rad(-z3).view(-1, 1) - cos_t = torch.cos(theta) - sin_t = torch.sin(theta) - - rays_x_rot1 = cos_t * rays_x - sin_t * rays_y - rays_y_rot1 = sin_t * rays_x + cos_t * rays_y - rays_z_rot1 = rays_z - - theta = torch.deg2rad(x).view(-1, 1) - cos_t = torch.cos(theta) - sin_t = torch.sin(theta) - - rays_x_rot2 = rays_x_rot1 - rays_y_rot2 = cos_t * rays_y_rot1 - sin_t * rays_z_rot1 - rays_z_rot2 = sin_t * rays_y_rot1 + cos_t * rays_z_rot1 - - theta = torch.deg2rad(-z1).view(-1, 1) - cos_t = torch.cos(theta) - sin_t = torch.sin(theta) - - rays_x_final = cos_t * rays_x_rot2 - sin_t * rays_y_rot2 - rays_y_final = sin_t * rays_x_rot2 + cos_t * rays_y_rot2 - rays_z_final = rays_z_rot2 - - transformed_rays = torch.stack([rays_x_final, rays_y_final, rays_z_final], dim=2) + shifted = torch.stack( + [rays[:, :, 0] - shift_x_norm, rays[:, :, 1] - shift_y_norm, rays[:, :, 2]], + dim=2, + ) - return transformed_rays + # Compose the three Euler rotations Rz(-z1) @ Rx(x) @ Rz(-z3) into a single + # (B, 3, 3) matrix and apply it with one batched matmul, instead of nine + # elementwise passes over the full (B, S) ray tensors. + a = torch.deg2rad(-z3).view(-1) + b = torch.deg2rad(x).view(-1) + g = torch.deg2rad(-z1).view(-1) + zero = torch.zeros_like(a) + one = torch.ones_like(a) + + cos_a, sin_a = torch.cos(a), torch.sin(a) + cos_b, sin_b = torch.cos(b), torch.sin(b) + cos_g, sin_g = torch.cos(g), torch.sin(g) + + rot_a = torch.stack( + [cos_a, -sin_a, zero, sin_a, cos_a, zero, zero, zero, one], dim=-1 + ).view(-1, 3, 3) + rot_b = torch.stack( + [one, zero, zero, zero, cos_b, -sin_b, zero, sin_b, cos_b], dim=-1 + ).view(-1, 3, 3) + rot_g = torch.stack( + [cos_g, -sin_g, zero, sin_g, cos_g, zero, zero, zero, one], dim=-1 + ).view(-1, 3, 3) + + rot = rot_g @ rot_b @ rot_a # (B, 3, 3) + + return shifted @ rot.transpose(1, 2) @staticmethod @torch.compile(mode="reduce-overhead") @@ -614,10 +614,13 @@ def __getitem__( pixel_i = remaining // self.tilt_stack.shape[2] pixel_j = remaining % self.tilt_stack.shape[2] + # Plain ints for the index fields: default_collate builds one int64 tensor per + # batch either way, but wrapping each index in torch.tensor() here allocates + # three scalar tensors per item on the dataloader hot path. return { - "projection_idx": torch.tensor(projection_idx), - "pixel_i": torch.tensor(pixel_i), - "pixel_j": torch.tensor(pixel_j), + "projection_idx": projection_idx, + "pixel_i": pixel_i, + "pixel_j": pixel_j, "phi": self.tilt_angles[projection_idx], # tensor "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j], # tensor } From 015299232c48ae8a9193187fe315ab09019c8f96 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 03:30:05 -0700 Subject: [PATCH 07/19] Fix no-op SIRT inline alignment and multi-angle rotation operators - TomographyConventional._reconstruction_epoch wrote the aligned measurement into proj_forward, which radon_torch overwrites immediately after, while the error term keeps reading the original tilt stack -- inline_alignment was a no-op. The aligned image is now persisted into the tilt stack (the tensor the error term reads), and the measurement is transposed to match the slice/detector orientation of the forward projection it is correlated with. - differentiable_rotz/rotx_vectorized raised for more than one angle: the per-slice vmap built affine_grid with a (T, 2, 3) matrix against a slice batch of 1. A rotation about an axis applies the same 2-D transform to every slice along it, so that axis can ride along as grid_sample channels in a single call -- which both fixes multi-angle/per-volume batching and removes the vmap. Scalar-angle outputs are unchanged (verified to 1e-6). --- src/quantem/tomography/tomography.py | 11 +++- src/quantem/tomography/utils.py | 64 +++++++++---------- .../test_tomography_conventional.py | 21 ++++++ tests/tomography/test_utils.py | 24 +++++++ 4 files changed, 86 insertions(+), 34 deletions(-) diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index f86095ad..db5656b4 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -566,7 +566,10 @@ def _reconstruction_epoch( if inline_alignment: for ind in range(len(self.dset.tilt_angles)): im_proj = proj_forward[:, ind, :] - im_meas = self.dset.forward(ind).target # type: ignore + # proj_forward rows are volume slices (the tilt axis), i.e. the + # transpose of the stored tilt image -- the same orientation the + # error term below compares against. + im_meas = self.dset.forward(ind).target.T # type: ignore shift = torch_phase_cross_correlation(im_proj, im_meas) if torch.linalg.norm(shift) <= 32: shifted = torch.fft.ifft2( @@ -585,7 +588,11 @@ def _reconstruction_epoch( ) ).real - proj_forward[:, ind, :] = shifted + # Persist the aligned measurement in the tilt stack: the error + # term reads the stack, and proj_forward is overwritten by + # radon_torch below, so writing the aligned image there + # silently discarded the alignment. + self.dset.tilt_stack[ind] = shifted.T if mode == "sirt" or mode == "fbp": proj_forward = radon_torch( diff --git a/src/quantem/tomography/utils.py b/src/quantem/tomography/utils.py index 08093925..7a592bf1 100644 --- a/src/quantem/tomography/utils.py +++ b/src/quantem/tomography/utils.py @@ -19,54 +19,54 @@ def rot_ZXZ(mags, z1, x, z3, device, mode="bilinear"): return curr_mags -def differentiable_rotz_vectorized(mags, theta, mode="bilinear"): - _, dimz, dimy, dimx = mags.shape - - if theta.dim() == 0: - theta = theta.unsqueeze(0) - +def _rot2d_affine_matrix(theta: torch.Tensor, batch: int) -> torch.Tensor: + """(B, 2, 3) in-plane rotation matrices, broadcasting a single angle over the batch.""" theta_rad = torch.deg2rad(theta) - cos_t, sin_t = torch.cos(theta_rad), torch.sin(theta_rad) affine_matrix = torch.stack( [cos_t, -sin_t, torch.zeros_like(theta), sin_t, cos_t, torch.zeros_like(theta)], dim=-1 ).view(-1, 2, 3) - - mags = mags.permute(1, 0, 2, 3) - - def transform_slice(mag_slice): - grid = F.affine_grid(affine_matrix, mag_slice.unsqueeze(0).shape, align_corners=False) - return F.grid_sample(mag_slice.unsqueeze(0), grid, mode=mode, align_corners=False).squeeze( - 0 + if affine_matrix.shape[0] == 1 and batch > 1: + affine_matrix = affine_matrix.expand(batch, 2, 3) + elif affine_matrix.shape[0] != batch and batch != 1: + raise ValueError( + f"Got {affine_matrix.shape[0]} angles for a batch of {batch} volumes; " + "pass one angle, or one angle per volume." ) + return affine_matrix - rotated_mags = torch.vmap(transform_slice)(mags) - return rotated_mags.permute(1, 0, 2, 3) - -def differentiable_rotx_vectorized(mags, theta, mode="bilinear"): - _, dimz, dimy, dimx = mags.shape +def differentiable_rotz_vectorized(mags, theta, mode="bilinear"): + B, dimz, dimy, dimx = mags.shape if theta.dim() == 0: theta = theta.unsqueeze(0) - theta_rad = torch.deg2rad(theta) + # A rotation about z applies the same 2-D transform to every z-slice, so the + # z-axis rides along as grid_sample channels in a single call. (The previous + # per-slice vmap also broke for more than one angle, because affine_grid + # requires the matrix batch to match the slice batch of 1.) + affine_matrix = _rot2d_affine_matrix(theta, B) + if affine_matrix.shape[0] > B: # one volume, many angles + mags = mags.expand(affine_matrix.shape[0], dimz, dimy, dimx) + grid = F.affine_grid(affine_matrix, mags.shape, align_corners=False) + return F.grid_sample(mags, grid, mode=mode, align_corners=False) - cos_t, sin_t = torch.cos(theta_rad), torch.sin(theta_rad) - affine_matrix = torch.stack( - [cos_t, -sin_t, torch.zeros_like(theta), sin_t, cos_t, torch.zeros_like(theta)], dim=-1 - ).view(-1, 2, 3) - mags = mags.permute(3, 0, 1, 2) +def differentiable_rotx_vectorized(mags, theta, mode="bilinear"): + B, dimz, dimy, dimx = mags.shape - def transform_slice(mag_slice): - grid = F.affine_grid(affine_matrix, mag_slice.unsqueeze(0).shape, align_corners=False) - return F.grid_sample(mag_slice.unsqueeze(0), grid, mode=mode, align_corners=False).squeeze( - 0 - ) + if theta.dim() == 0: + theta = theta.unsqueeze(0) - rotated_mags = torch.vmap(transform_slice)(mags) - return rotated_mags.permute(1, 2, 3, 0) + # Same trick as rotz with the x-axis as the channel dim: rotate in (z, y). + affine_matrix = _rot2d_affine_matrix(theta, B) + mags = mags.permute(0, 3, 1, 2) # (B, X, Z, Y) + if affine_matrix.shape[0] > B: # one volume, many angles + mags = mags.expand(affine_matrix.shape[0], dimx, dimz, dimy) + grid = F.affine_grid(affine_matrix, mags.shape, align_corners=False) + rotated = F.grid_sample(mags, grid, mode=mode, align_corners=False) + return rotated.permute(0, 2, 3, 1) # back to (B, Z, Y, X) def tv_loss_1d(x: torch.Tensor, reduction: str = "mean") -> torch.Tensor: diff --git a/tests/tomography/test_tomography_conventional.py b/tests/tomography/test_tomography_conventional.py index a7cfb5e0..85e3daf2 100644 --- a/tests/tomography/test_tomography_conventional.py +++ b/tests/tomography/test_tomography_conventional.py @@ -61,6 +61,27 @@ def test_obj_constraints_accepts_dict(self, phantom_volume, tilt_series, tilt_an assert tomo.num_epochs == 2 +class TestInlineAlignment: + def test_alignment_corrects_misaligned_projection( + self, phantom_volume, tilt_series, tilt_angles + ): + """Regression: the aligned measurement was written into proj_forward, which + radon_torch overwrites immediately, so inline_alignment was a no-op.""" + n = phantom_volume.shape[0] + shifted_series = tilt_series.copy() + shifted_series[4] = np.roll(shifted_series[4], shift=3, axis=1) + + reference = _build(tilt_series, tilt_angles, n).dset.tilt_stack[4] + tomo = _build(shifted_series, tilt_angles, n) + misaligned_before = tomo.dset.tilt_stack[4].clone() + tomo.reconstruct(num_iter=3, mode="sirt", inline_alignment=True) + aligned_after = tomo.dset.tilt_stack[4] + + err_before = float(((misaligned_before - reference) ** 2).mean()) + err_after = float(((aligned_after - reference) ** 2).mean()) + assert err_after < err_before # the stack was actually re-aligned + + class TestFBP: def test_fbp_runs_single_epoch(self, phantom_volume, tilt_series, tilt_angles): n = phantom_volume.shape[0] diff --git a/tests/tomography/test_utils.py b/tests/tomography/test_utils.py index 1fa6a47c..233e6bea 100644 --- a/tests/tomography/test_utils.py +++ b/tests/tomography/test_utils.py @@ -5,6 +5,7 @@ import torch from quantem.tomography.utils import ( + differentiable_rotx_vectorized, differentiable_rotz_vectorized, rot_ZXZ, tv_loss_1d, @@ -76,3 +77,26 @@ def test_gradient_flows_through_rotation(self): out.sum().backward() assert vol.grad is not None assert torch.isfinite(vol.grad).all() + + @pytest.mark.parametrize( + "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized] + ) + def test_multi_angle_matches_per_angle_calls(self, rot_fn): + # Regression: the per-slice vmap version raised for more than one angle + # (affine_grid requires the matrix batch to match the slice batch of 1). + vol = _block_volume() + angles = torch.tensor([10.0, -25.0, 60.0]) + multi = rot_fn(vol, angles) + per_angle = torch.cat([rot_fn(vol, a) for a in angles]) + assert multi.shape == (3, *vol.shape[1:]) + assert torch.allclose(multi, per_angle, atol=1e-6) + + @pytest.mark.parametrize( + "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized] + ) + def test_per_volume_angles(self, rot_fn): + vols = torch.cat([_block_volume(), _block_volume().flip(-1), _block_volume().flip(-2)]) + angles = torch.tensor([10.0, -25.0, 60.0]) + batched = rot_fn(vols, angles) + per_volume = torch.cat([rot_fn(vols[i : i + 1], angles[i]) for i in range(3)]) + assert torch.allclose(batched, per_volume, atol=1e-6) From 6a88afb260d044021e3633dfda2a949b91ef3800 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 04:26:56 -0700 Subject: [PATCH 08/19] Make scheduler params() pure and fix gradient-detaching angle wrapping in rot_ZXZ - Every SchedulerParams dataclass (Plateau, Exponential, Cyclic, Linear, CosineAnnealing) mutated its own fields inside params(): the first call permanently baked derived values (min_lr, gamma, base/max_lr, total_iters, T_max) into the instance, so a config shared between the object and pose optimizers -- or reused across reconstruct() calls with a different num_iter -- silently kept the first call's values. Derived values are now computed locally; explicitly-set fields still take precedence. - rot_ZXZ re-wrapped all three Euler angles with torch.tensor() whenever any one of them was a non-tensor, copying and detaching tensor angles and silently cutting gradient flow (plus a UserWarning). Each angle is now converted independently, leaving tensors untouched. --- src/quantem/core/ml/optimizer_mixin.py | 46 ++++++++++---------- src/quantem/tomography/utils.py | 11 +++-- tests/ml/test_optimizermixin.py | 58 +++++++++++++++++++++++--- tests/tomography/test_utils.py | 12 ++++++ 4 files changed, 96 insertions(+), 31 deletions(-) diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py index ece31dff..af125391 100644 --- a/src/quantem/core/ml/optimizer_mixin.py +++ b/src/quantem/core/ml/optimizer_mixin.py @@ -294,14 +294,15 @@ class Plateau: _name: str = "plateau" def params(self, base_LR: float, num_iter: int | None = None) -> dict: - if self.min_lr is None: - self.min_lr = self.min_lr_factor * base_LR + # Derived values stay local: params() must not mutate the dataclass, or a + # shared instance bakes in the first optimizer's base LR for every later one. + min_lr = self.min_lr if self.min_lr is not None else self.min_lr_factor * base_LR return { "mode": self.mode, "factor": self.factor, "patience": self.patience, "threshold": self.threshold, - "min_lr": self.min_lr, + "min_lr": min_lr, "cooldown": self.cooldown, } @@ -332,13 +333,12 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict: if effective_num_iter is None: raise ValueError("num_iter must be set if num_iter is not provided") - self.num_iter = effective_num_iter - + gamma = self.gamma if self.factor is not None: - self.gamma = self.factor ** (1.0 / effective_num_iter) + gamma = self.factor ** (1.0 / effective_num_iter) return { - "gamma": self.gamma, + "gamma": gamma, } @dataclass @@ -385,13 +385,11 @@ class Cyclic: _name: str = "cyclic" def params(self, base_LR: float, num_iter: int | None = None) -> dict: - if self.base_lr is None: - self.base_lr = self.base_lr_factor * base_LR - if self.max_lr is None: - self.max_lr = self.max_lr_factor * base_LR + base_lr = self.base_lr if self.base_lr is not None else self.base_lr_factor * base_LR + max_lr = self.max_lr if self.max_lr is not None else self.max_lr_factor * base_LR return { - "base_lr": self.base_lr, - "max_lr": self.max_lr, + "base_lr": base_lr, + "max_lr": max_lr, "step_size_up": self.step_size_up, "step_size_down": self.step_size_down, "mode": self.mode, @@ -426,12 +424,11 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict: raise ValueError( "total_iters must be set if num_iter is not provided" ) # Should never be reached - if self.total_iters is None: - self.total_iters = num_iter + total_iters = self.total_iters if self.total_iters is not None else num_iter return { "start_factor": self.start_factor, "end_factor": self.end_factor, - "total_iters": self.total_iters, + "total_iters": total_iters, } @dataclass @@ -459,10 +456,9 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict: raise ValueError( "T_max must be set if num_iter is not provided" ) # Should never be reached - if self.T_max is None: - self.T_max = num_iter + T_max = self.T_max if self.T_max is not None else num_iter return { - "T_max": self.T_max, + "T_max": T_max, "eta_min": self.eta_min, } @@ -572,7 +568,9 @@ def _normalize_optimizer_params( if isinstance(params, OptimizerParamsType): return {self.DEFAULT_OPTIMIZER_KEY: params} if not isinstance(params, dict): - raise TypeError(f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}") + raise TypeError( + f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}" + ) # Single optimizer as dict shorthand, e.g. {"name": "adam", "lr": 1e-3} if self._is_single_optimizer_dict(params): return {self.DEFAULT_OPTIMIZER_KEY: OptimizerParams.parse_dict(d=params)} @@ -597,7 +595,9 @@ def scheduler_params(self, params: SchedulerParamsType | dict): if isinstance(params, dict): params = SchedulerParams.parse_dict(d=params) if not isinstance(params, SchedulerParamsType): - raise TypeError(f"scheduler parameters must be a SchedulerParamsType, got {type(params)}") + raise TypeError( + f"scheduler parameters must be a SchedulerParamsType, got {type(params)}" + ) self._scheduler_params = params @abstractmethod @@ -688,7 +688,9 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer": raise NotImplementedError(f"Unknown optimizer type: {opt_params}") def set_scheduler( - self, scheduler_params: SchedulerParamsType | dict | None = None, num_iter: int | None = None + self, + scheduler_params: SchedulerParamsType | dict | None = None, + num_iter: int | None = None, ) -> None: """Set the scheduler for this model.""" if scheduler_params is not None: diff --git a/src/quantem/tomography/utils.py b/src/quantem/tomography/utils.py index 08093925..e86e8d13 100644 --- a/src/quantem/tomography/utils.py +++ b/src/quantem/tomography/utils.py @@ -5,10 +5,13 @@ def rot_ZXZ(mags, z1, x, z3, device, mode="bilinear"): - if not isinstance(x, torch.Tensor) or not isinstance(z1, torch.Tensor): - z1 = torch.tensor(z1, dtype=torch.float32, device=device) - x = torch.tensor(x, dtype=torch.float32, device=device) - z3 = torch.tensor(z3, dtype=torch.float32, device=device) + # Convert each angle independently: re-wrapping an existing tensor with + # torch.tensor() copies and detaches it, silently cutting gradient flow + # through tensor angles whenever any other angle is passed as a float. + z1, x, z3 = ( + a if isinstance(a, torch.Tensor) else torch.tensor(a, dtype=torch.float32, device=device) + for a in (z1, x, z3) + ) curr_mags = mags curr_mags = differentiable_rotz_vectorized(curr_mags, z1, mode) diff --git a/tests/ml/test_optimizermixin.py b/tests/ml/test_optimizermixin.py index d6347ce6..7dc6f4e7 100644 --- a/tests/ml/test_optimizermixin.py +++ b/tests/ml/test_optimizermixin.py @@ -356,9 +356,7 @@ def test_parse_default_name_is_none(self): from quantem.core.ml.optimizer_mixin import OptimizerMixin # noqa: E402 torch = pytest.importorskip("torch") if config.get("has_torch") else None -requires_torch = pytest.mark.skipif( - not config.get("has_torch"), reason="requires torch" -) +requires_torch = pytest.mark.skipif(not config.get("has_torch"), reason="requires torch") def _param(value=1.0): @@ -413,7 +411,10 @@ def test_none_optimizer_removes_without_touching_params(self): def test_multi_group_pplr_applies_per_group_lr(self): model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]}) model.set_optimizer( - {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)} + { + "descan": OptimizerParams.SGD(lr=1e-2), + "scan_positions": OptimizerParams.SGD(lr=1e-3), + } ) groups = model.optimizer.param_groups assert len(groups) == 2 @@ -445,9 +446,56 @@ def test_reset_optimizer_on_unconfigured_model_is_noop(self): def test_set_scheduler_base_lr_uses_max_group_lr(self): model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]}) model.set_optimizer( - {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)} + { + "descan": OptimizerParams.SGD(lr=1e-2), + "scan_positions": OptimizerParams.SGD(lr=1e-3), + } ) model.set_scheduler(SchedulerParams.Plateau(), num_iter=10) assert model.scheduler is not None # Plateau min_lr defaults to base_LR / 20, with base_LR = max group lr (1e-2) assert model.scheduler.min_lrs[0] == pytest.approx(1e-2 / 20) + + +class TestSchedulerParamsArePure: + """Regression: params() mutated the dataclass, so a shared instance baked in the + first call's derived values (min_lr, gamma, bounds, horizons) for every later + optimizer or reconstruct() call.""" + + def test_plateau_min_lr_follows_base_lr(self): + p = SchedulerParams.Plateau() + assert p.params(base_LR=1.0)["min_lr"] == pytest.approx(1 / 20) + assert p.params(base_LR=10.0)["min_lr"] == pytest.approx(10 / 20) + assert p.min_lr is None # instance untouched + + def test_exponential_gamma_follows_num_iter(self): + p = SchedulerParams.Exponential(factor=0.5) + g10 = p.params(base_LR=1.0, num_iter=10)["gamma"] + g100 = p.params(base_LR=1.0, num_iter=100)["gamma"] + assert g10 == pytest.approx(0.5 ** (1 / 10)) + assert g100 == pytest.approx(0.5 ** (1 / 100)) + assert p.num_iter is None and p.gamma == 0.9 + + def test_cyclic_bounds_follow_base_lr(self): + p = SchedulerParams.Cyclic() + first = p.params(base_LR=1.0) + second = p.params(base_LR=2.0) + assert second["base_lr"] == pytest.approx(2 * first["base_lr"]) + assert second["max_lr"] == pytest.approx(2 * first["max_lr"]) + assert p.base_lr is None and p.max_lr is None + + def test_linear_and_cosine_follow_num_iter(self): + lin = SchedulerParams.Linear() + assert lin.params(base_LR=1.0, num_iter=5)["total_iters"] == 5 + assert lin.params(base_LR=1.0, num_iter=50)["total_iters"] == 50 + assert lin.total_iters is None + cos = SchedulerParams.CosineAnnealing() + assert cos.params(base_LR=1.0, num_iter=5)["T_max"] == 5 + assert cos.params(base_LR=1.0, num_iter=50)["T_max"] == 50 + assert cos.T_max is None + + def test_explicit_values_still_win(self): + p = SchedulerParams.Plateau(min_lr=1e-6) + assert p.params(base_LR=1.0)["min_lr"] == 1e-6 + e = SchedulerParams.Exponential(gamma=0.8) + assert e.params(base_LR=1.0, num_iter=10)["gamma"] == 0.8 diff --git a/tests/tomography/test_utils.py b/tests/tomography/test_utils.py index 1fa6a47c..6994abbd 100644 --- a/tests/tomography/test_utils.py +++ b/tests/tomography/test_utils.py @@ -76,3 +76,15 @@ def test_gradient_flows_through_rotation(self): out.sum().backward() assert vol.grad is not None assert torch.isfinite(vol.grad).all() + + +class TestRotZXZGradients: + def test_grad_flows_with_mixed_float_and_tensor_angles(self): + """Regression: a non-tensor angle made rot_ZXZ re-wrap every angle with + torch.tensor(), detaching gradients through tensor angles.""" + vol = _block_volume() + x = torch.tensor(20.0, requires_grad=True) + out = rot_ZXZ(vol, 0.0, x, 0.0, device="cpu") + out.sum().backward() + assert x.grad is not None + assert torch.isfinite(x.grad) From 9fbe4e37297b457c0bbc8b479c926108b4488e49 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 05:23:37 -0700 Subject: [PATCH 09/19] Fix ignored winner-initialization seed in Siren and remove dead learnable_tilts setter - Siren/HSiren created and seeded a torch.Generator for the winner initialization but drew the perturbation with torch.randn_like, which ignores generators -- so winner_initialization=72 (used by TomographyLiteINR) silently had no effect on reproducibility and every seed produced the same global-RNG noise. Draw from torch.randn with the seeded generator instead. - TomographyDatasetBase.learnable_tilts had a setter that wrote a private attribute the getter never read, so assignments appeared to succeed while silently doing nothing. The value is derived from the tilt series; the setter is removed so assignment now raises instead of lying. --- src/quantem/core/ml/inr.py | 24 +++++++------- src/quantem/tomography/dataset_models.py | 7 ++--- tests/ml/test_inr.py | 40 ++++++++++++++++++++++++ tests/tomography/test_dataset_models.py | 9 ++++++ 4 files changed, 65 insertions(+), 15 deletions(-) create mode 100644 tests/ml/test_inr.py diff --git a/src/quantem/core/ml/inr.py b/src/quantem/core/ml/inr.py index 4c197344..f4801705 100644 --- a/src/quantem/core/ml/inr.py +++ b/src/quantem/core/ml/inr.py @@ -111,18 +111,20 @@ def _build(self) -> None: self.net = nn.Sequential(*net_list) if self.winner_initialization: - if type(self.winner_initialization) is int: - rng = torch.Generator() - rng.manual_seed(self.winner_initialization) - else: - rng = torch.Generator() - rng.manual_seed(42) + seed = self.winner_initialization if type(self.winner_initialization) is int else 42 + rng = torch.Generator() + rng.manual_seed(seed) + # torch.randn_like ignores generators, so the noise must come from + # torch.randn with the seeded generator -- otherwise the "winner" seed + # silently has no effect and the perturbation is not reproducible. with torch.no_grad(): - self.net[0].linear.weight += ( # type: ignore[reportAttributeAccessIssue] - torch.randn_like(self.net[0].linear.weight) * 5 / self.first_omega_0 # type:ignore - ) - self.net[1].linear.weight += ( # type: ignore[reportAttributeAccessIssue] - torch.randn_like(self.net[1].linear.weight) * 0.1 / self.hidden_omega_0 # type:ignore + w0 = self.net[0].linear.weight # type: ignore[reportAttributeAccessIssue] + w0 += torch.randn(w0.shape, generator=rng, dtype=w0.dtype) * 5 / self.first_omega_0 + w1 = self.net[1].linear.weight # type: ignore[reportAttributeAccessIssue] + w1 += ( + torch.randn(w1.shape, generator=rng, dtype=w1.dtype) + * 0.1 + / self.hidden_omega_0 ) def forward(self, coords: torch.Tensor) -> torch.Tensor: diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 2255636c..d328c4b9 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -302,12 +302,11 @@ def reference_tilt_idx(self, reference_tilt_idx: int): @property def learnable_tilts(self) -> int: + # Derived from the tilt series (all tilts minus the fixed reference); there is + # deliberately no setter -- the old one wrote a private attribute this getter + # never read, so assignments appeared to succeed while doing nothing. return self.tilt_angles.shape[0] - 1 - @learnable_tilts.setter - def learnable_tilts(self, learnable_tilts: int): - self._learnable_tilts = learnable_tilts - @property def z1_params(self) -> torch.nn.Parameter: return self._z1_params diff --git a/tests/ml/test_inr.py b/tests/ml/test_inr.py new file mode 100644 index 00000000..c78bc1af --- /dev/null +++ b/tests/ml/test_inr.py @@ -0,0 +1,40 @@ +"""Tests for ``quantem.core.ml.inr`` Siren / HSiren construction.""" + +import torch + +from quantem.core.ml.inr import HSiren, Siren + + +def _first_layer_weight(winner_seed) -> torch.Tensor: + # Fix the global RNG so the base SineLayer init is identical across builds; + # only the winner-initialization perturbation varies with winner_seed. + torch.manual_seed(0) + model = HSiren(hidden_layers=1, hidden_features=8, winner_initialization=winner_seed) + return model.net[0].linear.weight.detach().clone() + + +class TestWinnerInitialization: + def test_same_seed_is_reproducible(self): + assert torch.equal(_first_layer_weight(72), _first_layer_weight(72)) + + def test_different_seeds_differ(self): + """Regression: the seeded torch.Generator was created but never used + (torch.randn_like ignores generators), so every winner seed produced + the same perturbation.""" + assert not torch.equal(_first_layer_weight(72), _first_layer_weight(73)) + + def test_true_uses_default_seed(self): + assert torch.equal(_first_layer_weight(True), _first_layer_weight(42)) + + def test_disabled_adds_no_perturbation(self): + torch.manual_seed(0) + base = Siren(hidden_layers=1, hidden_features=8, winner_initialization=False) + torch.manual_seed(0) + again = Siren(hidden_layers=1, hidden_features=8, winner_initialization=False) + assert torch.equal(base.net[0].linear.weight.detach(), again.net[0].linear.weight.detach()) + + +def test_forward_shape(): + model = HSiren(hidden_layers=1, hidden_features=8) + out = model(torch.rand(5, 3)) + assert out.shape == (5, 1) diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py index 04abd3b5..2299a8cb 100644 --- a/tests/tomography/test_dataset_models.py +++ b/tests/tomography/test_dataset_models.py @@ -192,3 +192,12 @@ def test_save_load_parameters_roundtrip(self, tmp_path): d2.to("cpu") d2.load_parameters(path) assert torch.allclose(d2.z1_params.detach(), d.z1_params.detach()) + + +def test_learnable_tilts_is_read_only(): + """Regression: the old setter wrote a private attribute the getter never read, + so assignment appeared to succeed while silently doing nothing.""" + d = TomographyPixDataset.from_data(_stack(), np.linspace(-60, 60, 5).astype(np.float32)) + with pytest.raises(AttributeError): + d.learnable_tilts = 3 + assert d.learnable_tilts == 4 From 329e300d432c8731d3a6659fc4c2982f701df7eb Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 06:24:25 -0700 Subject: [PATCH 10/19] Fix pose parameters being reset by to() and refuse anisotropic KPlanes resolutions - TomographyPixDataset.to() and TomographyINRDataset.to() rebuilt the z1/z3/ shift parameters from the initial-value buffers (_z1_angles etc.), which are never updated during training -- so any device move after training, including Tomography.from_file(...).to(device), silently reset every learned pose to zero. Both now go through a shared helper that moves the current parameter values once they exist and only uses the buffers on first materialization. - KPlanes/KPlanesTILTED allocate all three plane grids from one (3[, *T], C, res[1], res[0]) tensor, ignoring res[2]; an anisotropic resolution silently gave the XZ/YZ planes the wrong grid along z. Constructor now raises a clear ValueError instead (isotropic behavior unchanged; supporting true anisotropy needs per-plane parameters, which would change the checkpoint layout). --- src/quantem/core/ml/models/kplanes.py | 10 ++++++++++ src/quantem/tomography/dataset_models.py | 23 +++++++++++++++++------ tests/ml/test_kplanes.py | 20 ++++++++++++++++++++ tests/tomography/test_dataset_models.py | 18 ++++++++++++++++++ 4 files changed, 65 insertions(+), 6 deletions(-) create mode 100644 tests/ml/test_kplanes.py diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py index cdf55261..80fc7f81 100644 --- a/src/quantem/core/ml/models/kplanes.py +++ b/src/quantem/core/ml/models/kplanes.py @@ -169,6 +169,7 @@ class KPlanes(PPLR, TensorDecompositionModel): """ K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241 """ + def __init__( self, # Grid parameters @@ -199,6 +200,15 @@ def __init__( self.concat_features = concat_features self.density_activation = density_activation + # All three planes share one (3, C, res[1], res[0]) tensor, which ignores + # res[2]: an anisotropic resolution would silently give the XZ/YZ planes the + # wrong grid along z. Refuse it rather than misallocate. + if len(set(self.resolution)) != 1: + raise ValueError( + f"KPlanes currently requires an isotropic resolution, got {self.resolution}; " + "the plane grids are allocated as (res[1], res[0]) for all three axis pairs." + ) + self.grids = nn.ParameterList() self.feature_dim = 0 for res_mult in self.multiscale_res_multipliers: diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 2255636c..8c9b01dd 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -240,6 +240,21 @@ def get_optimization_parameters(self) -> dict[str, list[torch.Tensor]]: """ return {self.DEFAULT_OPTIMIZER_KEY: list(self.parameters())} + def _materialize_pose_parameters(self, device: str | torch.device): + """Create the learnable pose parameters, or move the existing ones. + + Once the parameters exist, their *current* (possibly trained) values are + moved; the initial-value buffers are only used on first materialization. + Rebuilding from the buffers on every call silently reset learned poses + whenever the dataset changed device (e.g. ``from_file(...).to(device)``). + """ + z1 = self._z1_params.data if hasattr(self, "_z1_params") else self._z1_angles + z3 = self._z3_params.data if hasattr(self, "_z3_params") else self._z3_angles + shifts = self._shifts_params.data if hasattr(self, "_shifts_params") else self._shifts + self._z1_params = nn.Parameter(z1.detach().to(device)) + self._z3_params = nn.Parameter(z3.detach().to(device)) + self._shifts_params = nn.Parameter(shifts.detach().to(device)) + # --- Forward pass --- @abstractmethod def forward( @@ -429,9 +444,7 @@ def to(self, device: str | torch.device): self.tilt_stack = self.tilt_stack.to(device) self.tilt_angles = self.tilt_angles.to(device) - self._z1_params = nn.Parameter(self._z1_angles.to(device)) - self._z3_params = nn.Parameter(self._z3_angles.to(device)) - self._shifts_params = nn.Parameter(self._shifts.to(device)) + self._materialize_pose_parameters(device) self._z1_ref = self._z1_ref.to(device) self._z3_ref = self._z3_ref.to(device) @@ -632,9 +645,7 @@ def __len__( return self.tilt_stack.shape[0] * N * N def to(self, device: torch.device | str): - self._z1_params = nn.Parameter(self._z1_angles.to(device)) - self._z3_params = nn.Parameter(self._z3_angles.to(device)) - self._shifts_params = nn.Parameter(self._shifts.to(device)) + self._materialize_pose_parameters(device) self._z1_ref = self._z1_ref.to(device) self._z3_ref = self._z3_ref.to(device) diff --git a/tests/ml/test_kplanes.py b/tests/ml/test_kplanes.py new file mode 100644 index 00000000..bdd865bc --- /dev/null +++ b/tests/ml/test_kplanes.py @@ -0,0 +1,20 @@ +"""Tests for ``quantem.core.ml.models.kplanes`` construction guards.""" + +import pytest + +from quantem.core.ml.models.kplanes import KPlanes, KPlanesTILTED + + +class TestResolutionValidation: + def test_isotropic_builds(self): + model = KPlanes(M_features=2, resolution=(16, 16, 16)) + assert model.grids[0].shape == (3, 2, 16, 16) + + def test_anisotropic_raises(self): + """Regression: the plane grids are allocated as (res[1], res[0]) for all three + axis pairs, ignoring res[2] -- anisotropic resolutions silently gave the + XZ/YZ planes the wrong grid along z instead of erroring.""" + with pytest.raises(ValueError, match="isotropic"): + KPlanes(M_features=2, resolution=(16, 16, 8)) + with pytest.raises(ValueError, match="isotropic"): + KPlanesTILTED(M_features=2, T=2, resolution=(16, 8, 16)) diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py index 04abd3b5..e5a8c472 100644 --- a/tests/tomography/test_dataset_models.py +++ b/tests/tomography/test_dataset_models.py @@ -192,3 +192,21 @@ def test_save_load_parameters_roundtrip(self, tmp_path): d2.to("cpu") d2.load_parameters(path) assert torch.allclose(d2.z1_params.detach(), d.z1_params.detach()) + + +@pytest.mark.parametrize("cls", [TomographyPixDataset, TomographyINRDataset]) +def test_to_preserves_trained_pose_parameters(cls): + """Regression: to() rebuilt the pose parameters from the initial-value buffers, + so any device move after training (e.g. from_file(...).to(device)) silently + reset the learned pose to zero.""" + d = cls.from_data(_stack(), np.linspace(-60, 60, 5, dtype="f4")) + d.to("cpu") + d.z1_params.data.fill_(0.37) + d.z3_params.data.fill_(-0.21) + d.shifts_params.data.fill_(1.5) + + d.to("cpu") + + assert torch.allclose(d.z1_params.detach(), torch.full_like(d.z1_params, 0.37)) + assert torch.allclose(d.z3_params.detach(), torch.full_like(d.z3_params, -0.21)) + assert torch.allclose(d.shifts_params.detach(), torch.full_like(d.shifts_params, 1.5)) From be7188e7c2f8e97011472c1a896d091064543091 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 19:03:31 -0700 Subject: [PATCH 11/19] Fix sampler parity test to match int-returning __getitem__ The DeviceBatchSampler parity test compared batch entries against __getitem__ with .to(dtype), but __getitem__ returns plain ints for the index keys since the dataloader overhead reduction. Coerce with torch.as_tensor before comparing. --- tests/tomography/test_device_batch_sampler.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py index 47bce96f..8e050752 100644 --- a/tests/tomography/test_device_batch_sampler.py +++ b/tests/tomography/test_device_batch_sampler.py @@ -27,9 +27,14 @@ def test_batches_match_getitem(): for batch in sampler: for k in range(len(batch["target_value"])): item = dset[seen + k] + # __getitem__ returns plain ints for the index keys (cheaper than + # 0-d tensors); coerce before comparing against the batched tensors. 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 + batch[key][k], + torch.as_tensor(item[key], dtype=batch[key].dtype), + rtol=0, + atol=0, ) seen += len(batch["target_value"]) From 62e871e36fba4de2d990289518a9a70835d6f082 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 19:03:31 -0700 Subject: [PATCH 12/19] Guard tilt-stack normalization against sparse data A tilt stack with more than 95% zero pixels has a zero 95th quantile, so the normalization divided by zero and produced inf/NaN targets that poison every parameter on the first backward. Fall back to the absolute maximum when the quantile is non-positive and raise a clear error for all-zero stacks; same guard for the pretrain dataset. Adds regression tests for both paths. --- src/quantem/tomography/dataset_models.py | 13 +++++++++++++ tests/tomography/test_dataset_models.py | 16 ++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index b6c2c4ab..8b6cb8ba 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -189,6 +189,12 @@ def __init__( if type(tilt_angles) is not torch.Tensor: tilt_angles = torch.from_numpy(tilt_angles) max_val = torch.quantile(tilt_stack, 0.95) + # A sparse stack (>95% zeros) has a zero 95th quantile; dividing by it + # would turn the targets into inf/NaN and poison the first backward. + if max_val <= 0: + max_val = tilt_stack.abs().max() + if max_val <= 0: + raise ValueError("tilt_stack is all zeros; cannot normalize.") # Tilt stack normalization tilt_stack = tilt_stack / max_val @@ -793,6 +799,13 @@ def __init__( else: data_quantile = torch.quantile(data, 0.95) + # Same guard as TomographyDatasetBase: a >95%-zero target has a zero + # 95th quantile and would normalize to inf/NaN. + if data_quantile <= 0: + data_quantile = data.abs().max() + if data_quantile <= 0: + raise ValueError("pretrain_target is all zeros; cannot normalize.") + data = data / data_quantile data = torch.permute(data, (0, 3, 2, 1)) # data = torch.flip(data, dims=(2,)) diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py index d8148770..268fa2c6 100644 --- a/tests/tomography/test_dataset_models.py +++ b/tests/tomography/test_dataset_models.py @@ -58,6 +58,20 @@ def test_normalised_by_95th_quantile(self): q95 = torch.quantile(d.tilt_stack, 0.95) assert torch.isclose(q95, torch.tensor(1.0), atol=1e-4) + def test_sparse_stack_normalisation_stays_finite(self): + # >95% zeros makes the 95th quantile 0; normalisation must not divide + # by it (inf/NaN targets poison every parameter on the first backward). + stack = np.zeros((5, 12, 12), dtype=np.float32) + stack[:, 5:7, 5:7] = 10.0 + d = TomographyPixDataset.from_data(stack, np.linspace(-60, 60, 5).astype(np.float32)) + assert torch.isfinite(d.tilt_stack).all() + assert d.tilt_stack.max() > 0 + + def test_all_zero_stack_raises(self): + stack = np.zeros((5, 12, 12), dtype=np.float32) + with pytest.raises(ValueError): + TomographyPixDataset.from_data(stack, np.linspace(-60, 60, 5).astype(np.float32)) + def test_reference_idx_and_learnable_tilts(self): # negated angles -> [40, 15, -10, -35, -60]; smallest |angle| is index 2. angles = np.linspace(-40, 60, 5).astype(np.float32) @@ -229,6 +243,8 @@ def test_to_preserves_trained_pose_parameters(cls): assert torch.allclose(d.z1_params.detach(), torch.full_like(d.z1_params, 0.37)) assert torch.allclose(d.z3_params.detach(), torch.full_like(d.z3_params, -0.21)) assert torch.allclose(d.shifts_params.detach(), torch.full_like(d.shifts_params, 1.5)) + + def test_learnable_tilts_is_read_only(): """Regression: the old setter wrote a private attribute the getter never read, so assignment appeared to succeed while silently doing nothing.""" From 6496aa662038e531f20768c2c218780f1a5a85bf Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 22:20:46 -0700 Subject: [PATCH 13/19] Fix TV gate so plane-only TV is applied for tensor-decomp models apply_soft_constraints keyed the TV term on tv_vol alone, so setting tv_plane > 0 with tv_vol = 0 silently dropped the plane TV (the only TV used in the K-Planes paper, Eq. 3). Gate on either weight, and skip each TV term whose weight is zero so plane-only runs avoid the volume TV's three extra forward passes. --- src/quantem/tomography/object_models.py | 8 +++++--- tests/tomography/test_object_models.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index dad92a39..db66d48d 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -921,7 +921,7 @@ def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor: soft_loss = torch.tensor( 0.0, device=ctx.pred.device if ctx.pred is not None else self.device ) - if self.constraints.tv_vol > 0: + if self.constraints.tv_vol > 0 or self.constraints.tv_plane > 0: assert ctx.coords is not None, "Coordinates must be provided for TV loss" assert ctx.pred is not None, "Prediction must be provided for TV loss" soft_loss += self.get_tv_loss(ctx) @@ -947,8 +947,10 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor: assert ctx.coords is not None, "Coordinates must be provided for TV loss" assert ctx.pred is not None, "Prediction must be provided for TV loss" tv_loss = torch.tensor(0.0, device=ctx.pred.device) - tv_loss += self._get_plane_tv_loss() - tv_loss += self.get_volume_tv_loss(ctx.coords) + if self.constraints.tv_plane > 0: + tv_loss = tv_loss + self._get_plane_tv_loss() + if self.constraints.tv_vol > 0: + tv_loss = tv_loss + self.get_volume_tv_loss(ctx.coords) return tv_loss def _get_plane_tv_loss(self) -> torch.Tensor: diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py index d58314bb..da01c7d8 100644 --- a/tests/tomography/test_object_models.py +++ b/tests/tomography/test_object_models.py @@ -283,6 +283,17 @@ def test_volume_tv_loss_scalar(self, torch_device): assert loss.ndim == 0 assert torch.isfinite(loss) + def test_soft_constraints_plane_tv_only(self, torch_device): + """Regression: tv_plane > 0 with tv_vol = 0 must still apply the plane TV. + The gate previously keyed on tv_vol alone, silently dropping plane-only TV.""" + obj = self._obj(torch_device) + obj.constraints.tv_plane = 0.1 + obj.constraints.tv_vol = 0.0 + coords = torch.rand(64, 3, device=torch_device) * 2 - 1 + ctx = ReconstructionContext(coords=coords, pred=torch.zeros(64, device=torch_device)) + loss = obj.apply_soft_constraints(ctx) + assert float(loss.detach()) > 0.0 + def test_normalize_optimizer_params_rejects_non_dict(self, torch_device): from quantem.core.ml.optimizer_mixin import OptimizerParams From 6ed9fdcd1dcf12022f2ae006543615c283ae573d Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 03:33:07 -0700 Subject: [PATCH 14/19] Hoist loop-invariant sampling grid in CP-TILTED interpolation and drop index tensors in interpolate_ms_features - interpolate_ms_features_cp_tilted rebuilt the (3T, 1, B, 2) sampling grid (reshape/permute plus a zeros_like and stack) inside the per-scale loop even though it only depends on the rotated points; build it once before the loop. - interpolate_ms_features (non-tilted) projected points onto the three planes with list-based advanced indexing, which allocates an index tensor on device (a host-to-device copy) three times per forward; use unbind/stack views, matching the TILTED variant. Both bitwise-identical to the previous implementations; ~1.3x each on the isolated forward (B=4096, 3 scales; CP with T=8). --- src/quantem/core/ml/models/kplanes.py | 51 +++++++++++++-------------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py index 93dad498..3bbde431 100644 --- a/src/quantem/core/ml/models/kplanes.py +++ b/src/quantem/core/ml/models/kplanes.py @@ -144,12 +144,15 @@ def interpolate_ms_features( pts: torch.Tensor, ms_grids: nn.ParameterList, ) -> torch.Tensor: - mat_mode = [[0, 1], [0, 2], [1, 2]] + # Plane axis layout: XY=(0,1), XZ=(0,2), YZ=(1,2). Stacking views avoids the + # per-call index-tensor allocation (a host-to-device copy) that list-based + # advanced indexing does three times per forward. + x, y, z = pts.unbind(-1) coord_plane = torch.stack( [ - pts[:, mat_mode[0]], - pts[:, mat_mode[1]], - pts[:, mat_mode[2]], + torch.stack((x, y), dim=-1), + torch.stack((x, z), dim=-1), + torch.stack((y, z), dim=-1), ] ).view(3, -1, 1, 2) @@ -618,33 +621,29 @@ def interpolate_ms_features_cp_tilted( # Rotate all points by all rotations: (T, B, 3) rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) + # For each transform t, we need three 1D samples: at x_t, y_t, z_t. + # Lay them out as (3T, B) coords, matching each line grid's first dim. + # Axis order per transform: x, y, z. + coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B) + + # grid_sample wants 4D input for 2D sampling: sample the (3T, C, 1, L) lines + # with 2D coords whose y is fixed at 0. The grid only depends on the points, + # so it is built once here rather than per scale inside the loop. + grid = torch.stack( + [ + coords_1d, # x + torch.zeros_like(coords_1d), # y + ], + dim=-1, + ).unsqueeze(1) # (3T, 1, B, 2) + per_scale_features = [] for line_coef in ms_grids: # line_coef: (3T, C, L) — three 1D feature lines per transform (x, y, z) - C, _ = line_coef.shape[1], line_coef.shape[2] - - # For each transform t, we need three 1D samples: at x_t, y_t, z_t. - # Lay them out as (3T, B) coords, matching line_coef's first dim. - # Axis order per transform: x, y, z. - coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B) - # coords_1d: (3T, B), each row is samples along one axis for one transform - - # grid_sample wants 4D input for 2D sampling, or we can use 1D via a - # (3T, C, 1, L) reshape and pass 2D coords with y fixed at 0. - # Simpler: use F.grid_sample with a 4D trick, or just do manual linear interp. - # Here's the grid_sample way: - line_coef_4d = line_coef.unsqueeze(2) # (3T, C, 1, L) - # grid: need (3T, Hout=1, Wout=B, 2), with x = coord, y = 0 - grid = torch.stack( - [ - coords_1d, # x - torch.zeros_like(coords_1d), # y - ], - dim=-1, - ).unsqueeze(1) # (3T, 1, B, 2) + C = line_coef.shape[1] sampled = F.grid_sample( - line_coef_4d, + line_coef.unsqueeze(2), # (3T, C, 1, L) grid, align_corners=True, mode="bilinear", From 2640ba6c1a97b9a48e02fd8358b32d2b3b5b3e54 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 04:31:08 -0700 Subject: [PATCH 15/19] Vectorize INR dataset batch fetching and stack epoch loss reductions - TomographyINRDataset now implements __getitems__: torch's DataLoader calls it with the whole batch of indices, replacing one Python __getitem__ call (plus per-item dict construction and collate stacking) per sample with a few tensor ops. setup_dataloader pairs datasets that define __getitems__ with a passthrough collate_fn (module-level so spawn workers can pickle it); the random_split Subset path delegates the batched form automatically. Batches are identical to default collate (values and dtypes). 39.7x faster batch fetching (4.16 -> 0.105 ms/batch, batch_size=1024, 60x256x256 stack). - The per-epoch loss reduction issued three all_reduce calls and three .item() host syncs; stack the three scalars into one tensor for a single all_reduce and a single .tolist() sync. --- src/quantem/core/ml/ddp.py | 17 ++++++++++++++++- src/quantem/tomography/dataset_models.py | 22 ++++++++++++++++++++++ src/quantem/tomography/tomography.py | 13 ++++++------- 3 files changed, 44 insertions(+), 8 deletions(-) diff --git a/src/quantem/core/ml/ddp.py b/src/quantem/core/ml/ddp.py index 18b7ccd7..6627a7ed 100644 --- a/src/quantem/core/ml/ddp.py +++ b/src/quantem/core/ml/ddp.py @@ -12,6 +12,14 @@ def worker_init_fn(worker_id): os.environ["CUDA_VISIBLE_DEVICES"] = "" +def _passthrough_collate(batch): + """Identity collate for datasets whose ``__getitems__`` already returns the batch. + + Module-level (not a lambda) so spawn-context dataloader workers can pickle it. + """ + return batch + + class DDPMixin: """ Class for setting up all distributed training. @@ -97,6 +105,12 @@ def setup_dataloader( val_sampler = None shuffle = True + # Datasets that implement __getitems__ return an already-collated batch from a + # single vectorized call; pair them with a passthrough collate. (Gate on the + # original dataset: random_split's Subset defines __getitems__ unconditionally + # but only delegates the batched form when the wrapped dataset has it.) + collate_fn = _passthrough_collate if hasattr(dataset, "__getitems__") else None + train_dataloader = DataLoader( train_dataset, # type: ignore[reportArgumentType] --> Torch datasets do not have a len method, but still works. batch_size=batch_size, @@ -108,6 +122,7 @@ def setup_dataloader( persistent_workers=persist, multiprocessing_context=mp_ctx, worker_init_fn=worker_init_fn, + collate_fn=collate_fn, ) if val_dataset: @@ -122,8 +137,8 @@ def setup_dataloader( persistent_workers=persist, multiprocessing_context=mp_ctx, worker_init_fn=worker_init_fn, + collate_fn=collate_fn, ) - val_dataloader = val_dataloader else: val_dataloader = None diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 8b6cb8ba..de618a45 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -733,6 +733,28 @@ def __getitem__( "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j], # tensor } + def __getitems__(self, indices: list[int]) -> dict: + """ + Vectorized batch fetch. Torch's DataLoader calls this with the whole batch of + indices when it exists, replacing one Python ``__getitem__`` call (plus dict and + collate work) per item with a few tensor ops. Returns the already-collated batch; + the dataloader must use a passthrough collate_fn (see ``DDPMixin.setup_dataloader``). + """ + idx = torch.as_tensor(indices, dtype=torch.long) + pixels_per_projection = self.tilt_stack.shape[1] * self.tilt_stack.shape[2] + projection_idx = idx // pixels_per_projection + remaining = idx % pixels_per_projection + pixel_i = remaining // self.tilt_stack.shape[2] + pixel_j = remaining % self.tilt_stack.shape[2] + + return { + "projection_idx": projection_idx, + "pixel_i": pixel_i, + "pixel_j": pixel_j, + "phi": self.tilt_angles[projection_idx], + "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j], + } + def __len__( self, ): diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index ad2e65ce..cc6d9832 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -258,14 +258,13 @@ def reconstruct( ) prev_R = R_now.clone() + # One stacked all_reduce and one host sync instead of three of each. + losses = torch.stack([total_loss, consistency_loss, epoch_soft_constraint_loss]) if self.world_size > 1: - dist.all_reduce(total_loss, dist.ReduceOp.AVG) - dist.all_reduce(consistency_loss, dist.ReduceOp.AVG) - dist.all_reduce(epoch_soft_constraint_loss, dist.ReduceOp.AVG) - - total_loss = total_loss.item() / len(self.dataloader) - consistency_loss = consistency_loss.item() / len(self.dataloader) - epoch_soft_constraint_loss = epoch_soft_constraint_loss.item() / len(self.dataloader) + dist.all_reduce(losses, dist.ReduceOp.AVG) + total_loss, consistency_loss, epoch_soft_constraint_loss = ( + losses / len(self.dataloader) + ).tolist() self.step_schedulers(loss=total_loss) # TODO: Maybe reorganize the losses so that the order makes sense lol. From e42fc9668a9c90dfb5053fa12daf4b92bc494e07 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 05:27:11 -0700 Subject: [PATCH 16/19] Enable fused Adam/AdamW on CUDA and drop redundant epoch metrics reduction - _build_optimizer now passes fused=True to Adam/AdamW when every parameter is on a CUDA device: the whole optimizer step runs in one fused kernel, 2.0x faster than the default foreach path on KPlanesTILTED-sized grids (1.24 -> 0.61 ms/step; ~9% end-to-end per training iteration, where opt.step was 19% of the profile). Same update rule; parameters agree to ~2e-6 after 5 steps (kernel-order float effects). - The epoch losses were re-wrapped in a tensor, all_reduced a second time (idempotent on already rank-averaged values), and synced to host again right after the stacked reduction; the redundant block is removed. --- src/quantem/core/ml/optimizer_mixin.py | 8 ++++++-- src/quantem/tomography/tomography.py | 12 +++--------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py index af125391..bf54b7ea 100644 --- a/src/quantem/core/ml/optimizer_mixin.py +++ b/src/quantem/core/ml/optimizer_mixin.py @@ -672,11 +672,15 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer": so each group's ``lr`` etc. overrides the optimizer-level default. ``NoneOptimizer`` must have been filtered out by the caller. """ + # Fused Adam/AdamW runs the whole step in one CUDA kernel (~2x faster than the + # default foreach path on large grids, same update rule); only valid when every + # parameter lives on a CUDA device. + fused = all(p.is_cuda for group in param_groups for p in group["params"]) match opt_params: case OptimizerParams.Adam(): - return torch.optim.Adam(param_groups) + return torch.optim.Adam(param_groups, fused=fused) case OptimizerParams.AdamW(): - return torch.optim.AdamW(param_groups) + return torch.optim.AdamW(param_groups, fused=fused) case OptimizerParams.SGD(): return torch.optim.SGD(param_groups) case OptimizerParams.NoneOptimizer(): diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index cc6d9832..54d8426f 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -309,15 +309,9 @@ def reconstruct( avg_val_loss = val_loss.item() / len(self.val_dataloader) - metrics = torch.tensor( - [total_loss, consistency_loss, epoch_soft_constraint_loss], device=self.device - ) - - if self.world_size > 1: - dist.all_reduce(metrics, dist.ReduceOp.AVG) - - total_loss, consistency_loss, epoch_soft_constraint_loss = metrics.tolist() - + # The three losses were already rank-averaged (and batch-normalized) right + # after the batch loop; re-reducing identical values here was a redundant + # all_reduce plus an extra host sync per epoch. pbar.set_description( f"Reconstruction | Loss: {total_loss:.5e}, Consistency Loss: {consistency_loss:.5e}, Soft Constraint Loss: {epoch_soft_constraint_loss:.5e}" ) From 8b77d4a575dcc6f58a61e9623e2282f9f169e3b6 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 23:22:57 -0700 Subject: [PATCH 17/19] Revert "Merge pull request #12 from cedriclim1/perf/inr-batch-vectorize" This reverts commit d8454511e0e176bdc58005ea71ca38f30e63a8f2, reversing changes made to f8d9fb96a203960e646db95b64d395b6c50ac1ac. --- src/quantem/core/ml/ddp.py | 17 +--------------- src/quantem/core/ml/optimizer_mixin.py | 8 ++------ src/quantem/tomography/dataset_models.py | 22 --------------------- src/quantem/tomography/tomography.py | 25 +++++++++++++++--------- 4 files changed, 19 insertions(+), 53 deletions(-) diff --git a/src/quantem/core/ml/ddp.py b/src/quantem/core/ml/ddp.py index 6627a7ed..18b7ccd7 100644 --- a/src/quantem/core/ml/ddp.py +++ b/src/quantem/core/ml/ddp.py @@ -12,14 +12,6 @@ def worker_init_fn(worker_id): os.environ["CUDA_VISIBLE_DEVICES"] = "" -def _passthrough_collate(batch): - """Identity collate for datasets whose ``__getitems__`` already returns the batch. - - Module-level (not a lambda) so spawn-context dataloader workers can pickle it. - """ - return batch - - class DDPMixin: """ Class for setting up all distributed training. @@ -105,12 +97,6 @@ def setup_dataloader( val_sampler = None shuffle = True - # Datasets that implement __getitems__ return an already-collated batch from a - # single vectorized call; pair them with a passthrough collate. (Gate on the - # original dataset: random_split's Subset defines __getitems__ unconditionally - # but only delegates the batched form when the wrapped dataset has it.) - collate_fn = _passthrough_collate if hasattr(dataset, "__getitems__") else None - train_dataloader = DataLoader( train_dataset, # type: ignore[reportArgumentType] --> Torch datasets do not have a len method, but still works. batch_size=batch_size, @@ -122,7 +108,6 @@ def setup_dataloader( persistent_workers=persist, multiprocessing_context=mp_ctx, worker_init_fn=worker_init_fn, - collate_fn=collate_fn, ) if val_dataset: @@ -137,8 +122,8 @@ def setup_dataloader( persistent_workers=persist, multiprocessing_context=mp_ctx, worker_init_fn=worker_init_fn, - collate_fn=collate_fn, ) + val_dataloader = val_dataloader else: val_dataloader = None diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py index bf54b7ea..af125391 100644 --- a/src/quantem/core/ml/optimizer_mixin.py +++ b/src/quantem/core/ml/optimizer_mixin.py @@ -672,15 +672,11 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer": so each group's ``lr`` etc. overrides the optimizer-level default. ``NoneOptimizer`` must have been filtered out by the caller. """ - # Fused Adam/AdamW runs the whole step in one CUDA kernel (~2x faster than the - # default foreach path on large grids, same update rule); only valid when every - # parameter lives on a CUDA device. - fused = all(p.is_cuda for group in param_groups for p in group["params"]) match opt_params: case OptimizerParams.Adam(): - return torch.optim.Adam(param_groups, fused=fused) + return torch.optim.Adam(param_groups) case OptimizerParams.AdamW(): - return torch.optim.AdamW(param_groups, fused=fused) + return torch.optim.AdamW(param_groups) case OptimizerParams.SGD(): return torch.optim.SGD(param_groups) case OptimizerParams.NoneOptimizer(): diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index de618a45..8b6cb8ba 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -733,28 +733,6 @@ def __getitem__( "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j], # tensor } - def __getitems__(self, indices: list[int]) -> dict: - """ - Vectorized batch fetch. Torch's DataLoader calls this with the whole batch of - indices when it exists, replacing one Python ``__getitem__`` call (plus dict and - collate work) per item with a few tensor ops. Returns the already-collated batch; - the dataloader must use a passthrough collate_fn (see ``DDPMixin.setup_dataloader``). - """ - idx = torch.as_tensor(indices, dtype=torch.long) - pixels_per_projection = self.tilt_stack.shape[1] * self.tilt_stack.shape[2] - projection_idx = idx // pixels_per_projection - remaining = idx % pixels_per_projection - pixel_i = remaining // self.tilt_stack.shape[2] - pixel_j = remaining % self.tilt_stack.shape[2] - - return { - "projection_idx": projection_idx, - "pixel_i": pixel_i, - "pixel_j": pixel_j, - "phi": self.tilt_angles[projection_idx], - "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j], - } - def __len__( self, ): diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index 54d8426f..ad2e65ce 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -258,13 +258,14 @@ def reconstruct( ) prev_R = R_now.clone() - # One stacked all_reduce and one host sync instead of three of each. - losses = torch.stack([total_loss, consistency_loss, epoch_soft_constraint_loss]) if self.world_size > 1: - dist.all_reduce(losses, dist.ReduceOp.AVG) - total_loss, consistency_loss, epoch_soft_constraint_loss = ( - losses / len(self.dataloader) - ).tolist() + dist.all_reduce(total_loss, dist.ReduceOp.AVG) + dist.all_reduce(consistency_loss, dist.ReduceOp.AVG) + dist.all_reduce(epoch_soft_constraint_loss, dist.ReduceOp.AVG) + + total_loss = total_loss.item() / len(self.dataloader) + consistency_loss = consistency_loss.item() / len(self.dataloader) + epoch_soft_constraint_loss = epoch_soft_constraint_loss.item() / len(self.dataloader) self.step_schedulers(loss=total_loss) # TODO: Maybe reorganize the losses so that the order makes sense lol. @@ -309,9 +310,15 @@ def reconstruct( avg_val_loss = val_loss.item() / len(self.val_dataloader) - # The three losses were already rank-averaged (and batch-normalized) right - # after the batch loop; re-reducing identical values here was a redundant - # all_reduce plus an extra host sync per epoch. + metrics = torch.tensor( + [total_loss, consistency_loss, epoch_soft_constraint_loss], device=self.device + ) + + if self.world_size > 1: + dist.all_reduce(metrics, dist.ReduceOp.AVG) + + total_loss, consistency_loss, epoch_soft_constraint_loss = metrics.tolist() + pbar.set_description( f"Reconstruction | Loss: {total_loss:.5e}, Consistency Loss: {consistency_loss:.5e}, Soft Constraint Loss: {epoch_soft_constraint_loss:.5e}" ) From a582f9ef98932982e2c9b7193d130eb38a43e0cd Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Wed, 10 Jun 2026 23:22:57 -0700 Subject: [PATCH 18/19] Revert "Merge pull request #11 from cedriclim1/perf/cp-ms-grid-hoist" This reverts commit f8d9fb96a203960e646db95b64d395b6c50ac1ac, reversing changes made to 6496aa662038e531f20768c2c218780f1a5a85bf. --- src/quantem/core/ml/models/kplanes.py | 51 ++++++++++++++------------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py index 3bbde431..93dad498 100644 --- a/src/quantem/core/ml/models/kplanes.py +++ b/src/quantem/core/ml/models/kplanes.py @@ -144,15 +144,12 @@ def interpolate_ms_features( pts: torch.Tensor, ms_grids: nn.ParameterList, ) -> torch.Tensor: - # Plane axis layout: XY=(0,1), XZ=(0,2), YZ=(1,2). Stacking views avoids the - # per-call index-tensor allocation (a host-to-device copy) that list-based - # advanced indexing does three times per forward. - x, y, z = pts.unbind(-1) + mat_mode = [[0, 1], [0, 2], [1, 2]] coord_plane = torch.stack( [ - torch.stack((x, y), dim=-1), - torch.stack((x, z), dim=-1), - torch.stack((y, z), dim=-1), + pts[:, mat_mode[0]], + pts[:, mat_mode[1]], + pts[:, mat_mode[2]], ] ).view(3, -1, 1, 2) @@ -621,29 +618,33 @@ def interpolate_ms_features_cp_tilted( # Rotate all points by all rotations: (T, B, 3) rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) - # For each transform t, we need three 1D samples: at x_t, y_t, z_t. - # Lay them out as (3T, B) coords, matching each line grid's first dim. - # Axis order per transform: x, y, z. - coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B) - - # grid_sample wants 4D input for 2D sampling: sample the (3T, C, 1, L) lines - # with 2D coords whose y is fixed at 0. The grid only depends on the points, - # so it is built once here rather than per scale inside the loop. - grid = torch.stack( - [ - coords_1d, # x - torch.zeros_like(coords_1d), # y - ], - dim=-1, - ).unsqueeze(1) # (3T, 1, B, 2) - per_scale_features = [] for line_coef in ms_grids: # line_coef: (3T, C, L) — three 1D feature lines per transform (x, y, z) - C = line_coef.shape[1] + C, _ = line_coef.shape[1], line_coef.shape[2] + + # For each transform t, we need three 1D samples: at x_t, y_t, z_t. + # Lay them out as (3T, B) coords, matching line_coef's first dim. + # Axis order per transform: x, y, z. + coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B) + # coords_1d: (3T, B), each row is samples along one axis for one transform + + # grid_sample wants 4D input for 2D sampling, or we can use 1D via a + # (3T, C, 1, L) reshape and pass 2D coords with y fixed at 0. + # Simpler: use F.grid_sample with a 4D trick, or just do manual linear interp. + # Here's the grid_sample way: + line_coef_4d = line_coef.unsqueeze(2) # (3T, C, 1, L) + # grid: need (3T, Hout=1, Wout=B, 2), with x = coord, y = 0 + grid = torch.stack( + [ + coords_1d, # x + torch.zeros_like(coords_1d), # y + ], + dim=-1, + ).unsqueeze(1) # (3T, 1, B, 2) sampled = F.grid_sample( - line_coef.unsqueeze(2), # (3T, C, 1, L) + line_coef_4d, grid, align_corners=True, mode="bilinear", From 8fac79d9b8f08a92e9d0108fb2a04e495437bfd2 Mon Sep 17 00:00:00 2001 From: Cedric Lim Date: Wed, 24 Jun 2026 11:46:45 -0700 Subject: [PATCH 19/19] Fix DDP device placement and wrapper unwrapping in tomography First multi-node (world_size>1) run surfaced two DDP-only bugs: - TomographyBase.__init__ moved the dataset to the stale local `device` arg ("cuda:0") instead of self.device (cuda:local_rank set by setup_distributed), stranding pose params on cuda:0 while the model and sampler lived on cuda:local_rank -> device-mismatch in _apply_pose's batched matmul. - ObjectTensorDecomp._get_plane_tv_loss read .tilted/.T off the DDP wrapper instead of the inner module; ObjectINR.create_volume unwrapped with isinstance(nn.DataParallel) which DDP does not subclass. All three now go through self.device / the existing _unwrap() helper. Validated on a 16-rank (4x4 A100) run: both datasets train and rank 0 saves volumes with no device or attribute errors. Single-GPU behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/quantem/tomography/object_models.py | 8 ++++---- src/quantem/tomography/tomography_base.py | 6 +++++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index db66d48d..56e94740 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -778,7 +778,7 @@ def create_volume(self, return_vol: bool = False): coords_1d = torch.linspace(-1, 1, N) x, y, z = torch.meshgrid(coords_1d, coords_1d, coords_1d, indexing="ij") inputs = torch.stack([x, y, z], dim=-1).reshape(-1, 3) - model = self.model.module if isinstance(self.model, nn.DataParallel) else self.model + model = _unwrap(self.model) inference_batch_size = 5 * N * N total_samples = N**3 @@ -957,10 +957,10 @@ def _get_plane_tv_loss(self) -> torch.Tensor: """ Gets the total-variation across the planes. """ - is_tilted = self.model.tilted + model = _unwrap(self.model) + is_tilted = model.tilted per_level = [] - model = _unwrap(self.model) for p in model.grids: # p: (3*T, C, H, W) for TILTED, (3, C, H, W) for KPlanes dh = (p[:, :, 1:, :] - p[:, :, :-1, :]).pow(2).mean(dim=(1, 2, 3)) @@ -968,7 +968,7 @@ def _get_plane_tv_loss(self) -> torch.Tensor: per_plane = dh + dw # (3*T,) or (3,) if is_tilted: - T = self.model.T + T = model.T per_rotation = per_plane.view(T, 3).sum(dim=1) # sum 3 planes per rotation level_tv = per_rotation.mean() # avg across rotations else: diff --git a/src/quantem/tomography/tomography_base.py b/src/quantem/tomography/tomography_base.py index 419b5ede..2301b7be 100644 --- a/src/quantem/tomography/tomography_base.py +++ b/src/quantem/tomography/tomography_base.py @@ -58,7 +58,11 @@ def __init__( print("Setting up DDP for obj_model") self.dset = dset - self.dset.to(device) + # Use self.device (set by setup_distributed to cuda:local_rank under DDP), + # NOT the local `device` arg, which stays "cuda:0" on every rank and would + # strand the dataset's pose params on cuda:0 while the model/sampler live + # on cuda:local_rank (device-mismatch in _apply_pose's batched matmul). + self.dset.to(self.device) # --- Properties --- @property