From cc3a8503c7f7097c6c7b9a7e04a20562447cf10d Mon Sep 17 00:00:00 2001 From: ccasert Date: Tue, 18 Aug 2026 18:29:27 -0700 Subject: [PATCH] fix: preserve FP32 tomography coordinates under BF16 --- src/quantem/core/ml/models/kplanes.py | 11 +++- src/quantem/core/ml/models/so3params.py | 14 ++--- src/quantem/tomography/dataset_models.py | 67 ++++++++++++++---------- src/quantem/tomography/object_models.py | 15 ++++-- src/quantem/tomography/tomography.py | 14 ++--- 5 files changed, 72 insertions(+), 49 deletions(-) diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py index 84b83dea..cba4fc88 100644 --- a/src/quantem/core/ml/models/kplanes.py +++ b/src/quantem/core/ml/models/kplanes.py @@ -309,6 +309,13 @@ def resolution(self, resolution: Sequence[int]): # --------------------------------------------------------------------------- +def _rotate_tilted_points_fp32( + pts: torch.Tensor, rotation_matrices: torch.Tensor +) -> torch.Tensor: + with torch.autocast(device_type=pts.device.type, enabled=False): + return torch.einsum("tij,bj->tbi", rotation_matrices.float(), pts.float()) + + def interpolate_ms_features_tilted( pts: torch.Tensor, # (B, 3) ms_grids: nn.ParameterList, # each grid: (3*T, C, H, W) @@ -322,7 +329,7 @@ def interpolate_ms_features_tilted( B = pts.shape[0] # (T, B, 3) — rotate all points by all rotations at once - rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) + rotated = _rotate_tilted_points_fp32(pts, rotation_matrices) # 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. @@ -610,7 +617,7 @@ def interpolate_ms_features_cp_tilted( B = pts.shape[0] # Rotate all points by all rotations: (T, B, 3) - rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) + rotated = _rotate_tilted_points_fp32(pts, rotation_matrices) per_scale_features = [] for line_coef in ms_grids: diff --git a/src/quantem/core/ml/models/so3params.py b/src/quantem/core/ml/models/so3params.py index cd55d0ac..9a22a80a 100644 --- a/src/quantem/core/ml/models/so3params.py +++ b/src/quantem/core/ml/models/so3params.py @@ -174,11 +174,13 @@ def rotmat_to_r9(R: torch.Tensor) -> torch.Tensor: @staticmethod def r9_to_rotmat(M: torch.Tensor) -> torch.Tensor: """R9 (..., 3, 3) -> nearest SO(3) matrix via SVD+.""" - U, _, Vh = torch.linalg.svd(M) - d = torch.det(U @ Vh) - diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M.dtype) - diag[..., 2] = d - return U @ (diag.unsqueeze(-1) * Vh) + with torch.autocast(device_type=M.device.type, enabled=False): + M_work = M if M.dtype in (torch.float32, torch.float64) else M.float() + U, _, Vh = torch.linalg.svd(M_work) + d = torch.det(U @ Vh) + diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M_work.dtype) + diag[..., 2] = d + return U @ (diag.unsqueeze(-1) * Vh) def as_matrix(self) -> torch.Tensor: return self.r9_to_rotmat(self.M) @@ -189,4 +191,4 @@ def from_matrix(cls, R: torch.Tensor) -> "SO3ParamR9SVD": obj = cls(R.shape[0], init="identity") with torch.no_grad(): obj.M.copy_(cls.rotmat_to_r9(R)) - return obj \ No newline at end of file + return obj diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index a1a51bad..3d4ee028 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -508,32 +508,35 @@ def forward(self, dummy_input: Any = None): def get_coords( self, batch: dict[str, torch.Tensor], N: int, num_samples_per_ray: int ) -> torch.Tensor: - pixel_i = batch["pixel_i"].float().to(self.device, non_blocking=True) - pixel_j = batch["pixel_j"].float().to(self.device, non_blocking=True) - # target_values = batch["target_value"].to(self.device, non_blocking=True) - phis = batch["phi"].to(self.device, non_blocking=True) - projection_indices = batch["projection_idx"].to(self.device, non_blocking=True) - with torch.no_grad(): - batch_ray_coords = self.create_batch_rays(pixel_i, pixel_j, N, num_samples_per_ray) - - shifts, z1_params, z3_params = self.forward(None) - batch_shifts = torch.index_select(shifts, 0, projection_indices) - batch_z1 = torch.index_select(z1_params, 0, projection_indices) - batch_z3 = torch.index_select(z3_params, 0, projection_indices) - - transformed_rays = self.transform_batch_rays( - batch_ray_coords, - z1=batch_z1, - x=phis, - z3=batch_z3, - shifts=batch_shifts, - N=N, - sampling_rate=1.0, - ) - all_coords = transformed_rays.view(-1, 3) - - all_coords = all_coords.to(self.device, dtype=torch.float32, non_blocking=True) - return all_coords + with torch.autocast(device_type=self.device.type, enabled=False): + pixel_i = batch["pixel_i"].to( + self.device, dtype=torch.float32, non_blocking=True + ) + pixel_j = batch["pixel_j"].to( + self.device, dtype=torch.float32, non_blocking=True + ) + phis = batch["phi"].to(self.device, dtype=torch.float32, non_blocking=True) + projection_indices = batch["projection_idx"].to(self.device, non_blocking=True) + with torch.no_grad(): + batch_ray_coords = self.create_batch_rays( + pixel_i, pixel_j, N, num_samples_per_ray + ) + + shifts, z1_params, z3_params = self.forward(None) + batch_shifts = torch.index_select(shifts.float(), 0, projection_indices) + batch_z1 = torch.index_select(z1_params.float(), 0, projection_indices) + batch_z3 = torch.index_select(z3_params.float(), 0, projection_indices) + + transformed_rays = self.transform_batch_rays( + batch_ray_coords, + z1=batch_z1, + x=phis, + z3=batch_z3, + shifts=batch_shifts, + N=N, + sampling_rate=1.0, + ) + return transformed_rays.reshape(-1, 3) @staticmethod def create_batch_rays( @@ -542,9 +545,17 @@ def create_batch_rays( batch_size = len(pixel_i) x_coords = (pixel_j / (N - 1)) * 2 - 1 y_coords = (pixel_i / (N - 1)) * 2 - 1 - z_coords = torch.linspace(-1, 1, num_samples_per_ray, device=pixel_i.device) + z_coords = torch.linspace( + -1, 1, num_samples_per_ray, device=pixel_i.device, dtype=torch.float32 + ) - rays = torch.zeros(batch_size, num_samples_per_ray, 3, device=pixel_i.device) + rays = torch.zeros( + batch_size, + num_samples_per_ray, + 3, + device=pixel_i.device, + dtype=torch.float32, + ) rays[:, :, 0] = x_coords.unsqueeze(1) rays[:, :, 1] = y_coords.unsqueeze(1) diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index 817e7ed5..d1bbd044 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -20,6 +20,14 @@ from quantem.tomography.tomography_context import ReconstructionContext +def _tomography_autocast(device: torch.device | str, use_bfloat16: bool): + device = torch.device(device) + enabled = bool( + use_bfloat16 and device.type == "cuda" and torch.cuda.is_bf16_supported() + ) + return torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=enabled) + + class ObjConstraintParams: """ Namespace class for object reconstruction constraint dataclasses and parsing utilities. @@ -687,6 +695,7 @@ def pretrain( scheduler_params: dict | None = None, loss_fn: Callable | str = "l1", verbose: bool = True, + use_bfloat16: bool = True, ): """ Pretrain the INR model to fit target volume. @@ -717,6 +726,7 @@ def pretrain( num_iters=num_iters, loss_fn=loss_fn, verbose=verbose, + use_bfloat16=use_bfloat16, ) def _pretrain( @@ -724,6 +734,7 @@ def _pretrain( num_iters: int, loss_fn: Callable, verbose: bool, + use_bfloat16: bool, ): if self.optimizer is None: raise RuntimeError("Optimizer not set. Call set_optimizer() first.") @@ -741,9 +752,7 @@ def _pretrain( coords = batch["coords"].to(self.device, non_blocking=True) target = batch["target"].to(self.device, non_blocking=True) - with torch.autocast( - device_type=self.device.type, dtype=torch.bfloat16, enabled=True - ): + with _tomography_autocast(self.device, use_bfloat16): outputs = self.forward(coords) loss = loss_fn(outputs, target) diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index f86095ad..fa306049 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -27,6 +27,7 @@ ObjectINR, ObjectPixelated, ObjectTensorDecomp, + _tomography_autocast, ) from quantem.tomography.radon.radon import iradon_torch, radon_torch from quantem.tomography.tomography_base import TomographyBase @@ -84,6 +85,7 @@ def reconstruct( loss_func_kwargs: dict = {}, reset_dset: DatasetModelType | None = None, show_metrics: bool = False, + use_bfloat16: bool = True, ): """ This function should be able to handle both AD and INR-based tomography reconstruction methods. @@ -203,11 +205,7 @@ def reconstruct( for batch_idx, batch in enumerate(self.dataloader): self.zero_grad_all() - with torch.autocast( - device_type=self.device.type, - dtype=torch.bfloat16, - enabled=False, - ): + with _tomography_autocast(self.device, use_bfloat16): all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray) all_densities = self.obj_model.forward(all_coords) @@ -285,11 +283,7 @@ def reconstruct( val_loss = torch.tensor(0.0, device=self.device) for batch in self.val_dataloader: - with torch.autocast( - device_type=self.device.type, - dtype=torch.bfloat16, - enabled=True, - ): + with _tomography_autocast(self.device, use_bfloat16): all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray) all_densities = self.obj_model.forward(all_coords)