From 20360d87c7f37cac75874b0e674f7c578061a31e Mon Sep 17 00:00:00 2001 From: iback Date: Mon, 5 Oct 2026 13:13:54 +0000 Subject: [PATCH] perf: remove the per-sample loops and host stalls from the augmentations Every augmentation was read for work that scales with the batch rather than with the volume: Python loops over samples or regions, `.item()` reads inside those loops, and a few operations whose cost came from how they were spelled rather than from what they compute. On a [2, 1, 128, 128, 128] patch the 29 GPU transforms together drop from 86.8 ms to 56.2 ms, and the two heaviest -- the PALETTE pair -- from 43.6/18.0 ms to 28.3/9.5 ms. The three recurring causes: * **A one-channel `F.conv3d` is cuDNN's slowest path.** Folding the batch into the channel axis so `groups == B*C` selects ATen's depthwise kernel instead: 4.25 ms -> 0.12 ms for a 3x3x3 kernel over a 128^3 patch on an A40, and closer to the float64 reference as well, because the implicit-GEMM path it leaves runs in TF32. `kernels.depthwise_conv3d` is that reshape, and the convolution-based transforms now hand it the whole batch (with per-sample kernels zero-padded to a common width where RandConv draws different sizes). * **Reading a device scalar back to Python blocks the host.** `float(scales[b])`, `int(mask.sum())`, `bool(flips[b, axis])` and `torch.rand(1).item()` inside a per-sample loop each drain the queue; they are now single `.tolist()` transfers taken before the loop. * **`scatter_add_` into a handful of bins is an atomics contention test.** The region sums in PALETTE, the k-means cluster sums and the equalisation histogram all went through `bincount`, which accumulates in shared memory, or through a prefix sum. Beyond that: the bias-field polynomial is evaluated as a separable contraction rather than twenty full-volume terms; the EM/GMM E and M steps are two matrix multiplies instead of a dozen `(n, k)` passes; Voronoi sub-parcellation computes distances only for the cluster it is subdividing; the Scharr directions run as one convolution; `torch.quantile` per sample becomes one batched sort; and the constant kernels, the voxel-coordinate grid, the polynomial basis and the identity warp grid are cached on the patch shape instead of rebuilt per call. On the CPU backend the Laplacian uses its `3**d * delta - box` identity, the Scharr directions share one convolution, scipy's `norm.pdf` and its NumPy round trip are gone, and the redistribute dilation applies its structuring element three times rather than once pre-iterated: 110 -> 29 ms, 83 -> 33 ms, 26 -> 11 ms. Most transforms are bit-identical to before; `unit_tests/test_fast_paths.py` pins each reformulated primitive against the spelling it replaced. The exceptions are stated rather than hidden: * The convolution, bias-field and Scharr paths differ by float32 round-off (<= 2.7e-4 absolute, and the convolutions are *more* accurate than before). * PALETTE, the composed PALETTE, redistribute, the flip-axis guarantee, the bucketed random-choose pipelines and SynthSeg's background randomisation draw the same uniforms in a different order, so a seeded run no longer reproduces the previous release's output. Each was checked against a control that compares the old code with itself over a disjoint block of seeds: in every case the old-vs-new difference in mean, standard deviation and quantiles is smaller than that control. * `RandomHistogramEqualizationGPU` no longer raises out of `torch.histc` on a non-finite input; its own guard catches it and returns the channel untouched, which is what the guard was added for. The test that documented the old behaviour now pins the new one. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/cpu/contrast.py | 24 +- smauglab/transforms/cpu/fromSeg.py | 38 +- smauglab/transforms/gpu/contrast.py | 417 ++++++++++++------ smauglab/transforms/gpu/domain_transfer.py | 20 +- smauglab/transforms/gpu/fromSeg.py | 292 ++++++++---- smauglab/transforms/gpu/palette/base.py | 6 +- smauglab/transforms/gpu/palette/overlay.py | 5 +- .../transforms/gpu/palette/partitioners.py | 11 +- smauglab/transforms/gpu/palette/transform.py | 15 +- smauglab/transforms/gpu/spatial.py | 239 +++++----- smauglab/transforms/gpu/transforms_list.py | 29 +- smauglab/transforms/kernels.py | 143 +++++- smauglab/transforms/synthseg/functional.py | 175 ++++++-- unit_tests/test_fast_paths.py | 307 +++++++++++++ unit_tests/test_region_and_stats.py | 24 +- 15 files changed, 1280 insertions(+), 465 deletions(-) create mode 100644 unit_tests/test_fast_paths.py diff --git a/smauglab/transforms/cpu/contrast.py b/smauglab/transforms/cpu/contrast.py index c0ba94a..96ec012 100644 --- a/smauglab/transforms/cpu/contrast.py +++ b/smauglab/transforms/cpu/contrast.py @@ -8,7 +8,7 @@ from batchgeneratorsv2.transforms.intensity.gamma import GammaTransform from smauglab.registry import AugId, AugType, Backend, register -from smauglab.transforms.kernels import laplace_kernel, scharr_kernels +from smauglab.transforms.kernels import laplace_kernel, laplacian_response, scharr_kernels @register( @@ -85,15 +85,23 @@ def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: orig_mean = torch.mean(img[c]) orig_std = torch.std(img[c]) img_ = img[c].unsqueeze(0).unsqueeze(0) # adds temp batch and channel dim + spatial_dims = img_.dim() - 2 if params["kernel_type"] == "Laplace": - tot_ = apply_filter(img_, params["kernel"]) + # The separable identity rather than the dense kernel -- a quarter of + # the time for the same filter. See `kernels.laplacian_response`. + tot_ = laplacian_response(img_, spatial_dims) elif params["kernel_type"] == "Scharr": - tot_ = torch.zeros_like(img_) - for kernel in params["kernel"]: - if params["absolute"]: - tot_ += torch.abs(apply_filter(img_, kernel)) - else: - tot_ += apply_filter(img_, kernel) + # One convolution with every directional kernel in the output-channel + # axis, not one convolution per direction summed: bit-for-bit the same + # response, a third of the time. `params["kernel"]` still carries the + # kernels individually, so anything inspecting the parameters sees what + # it always did. + conv = F.conv3d if spatial_dims == 3 else F.conv2d + weight = torch.stack([k.to(img_) for k in params["kernel"]]).unsqueeze(1) + response = conv(img_, weight, padding="same") + if params["absolute"]: + response = response.abs() + tot_ = response.sum(dim=1, keepdim=True) img[c] = tot_[0, 0] if params["retain_stats"]: mean = torch.mean(img[c]) diff --git a/smauglab/transforms/cpu/fromSeg.py b/smauglab/transforms/cpu/fromSeg.py index 9e3e338..16547a5 100644 --- a/smauglab/transforms/cpu/fromSeg.py +++ b/smauglab/transforms/cpu/fromSeg.py @@ -1,9 +1,8 @@ -from functools import partial +import math import scipy.ndimage as ndi import torch from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform -from scipy.stats import norm from smauglab.registry import AugId, AugType, Backend, register @@ -86,11 +85,18 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) # Rank from the data, not hardcoded 3: scipy requires the structuring element # to match the input's rank, so a 2-D image raised # "structure rank must match input rank" here. - struct = ndi.iterate_structure(ndi.generate_binary_structure(l_mask_np.ndim, 1), 3) - l_mask_dilate_np = ndi.binary_dilation(l_mask_np, structure=struct) - - # Convert back to PyTorch - l_mask_dilate = torch.tensor(l_mask_dilate_np, device=device) + # + # Three passes of the 6-connected element, not one pass of + # `iterate_structure(..., 3)`. They give bit-identical masks -- that is what + # iterating a structure means -- but the iterated element is a 7x7x7 stencil + # evaluated at every voxel, where three small passes are 3 x 7 taps: 9 ms + # against 45 ms for a 96^3 mask. + struct = ndi.generate_binary_structure(l_mask_np.ndim, 1) + l_mask_dilate_np = ndi.binary_dilation(l_mask_np, structure=struct, iterations=3) + + # Convert back to PyTorch. `from_numpy` shares the buffer; `torch.tensor` + # copied it, and also warned about doing so. + l_mask_dilate = torch.from_numpy(l_mask_dilate_np).to(device) # Create mask of the dilated mask excluding the original mask l_mask_dilate_excl = l_mask_dilate & ~l_mask @@ -107,12 +113,11 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) torch.tensor([0.01], device=device), ) - redist = partial(norm.pdf, loc=l_mean.cpu().numpy(), scale=redist_std.cpu().numpy()) - + sign = 2 * torch.rand(1, device=device) - 1 if in_seg_bool: - to_add[l_mask] += torch.tensor(redist(img[l_mask].cpu().numpy()), device=device) * (2 * torch.rand(1, device=device) - 1) + to_add[l_mask] += _normal_pdf(img[l_mask], l_mean, redist_std) * sign else: - to_add += torch.tensor(redist(img.cpu().numpy()), device=device) * (2 * torch.rand(1, device=device) - 1) + to_add += _normal_pdf(img, l_mean, redist_std) * sign # Normalize to_add and apply it to the image. # @@ -139,6 +144,17 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) return img, seg +def _normal_pdf(x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: + """Gaussian density, in torch. + + This was `scipy.stats.norm.pdf` reached through a `functools.partial`, with the + volume copied out to NumPy and the result copied back for every label. The + formula is the same one `scipy` evaluates; what goes away is the frozen + distribution's argument handling and two full-volume round trips per label. + """ + return torch.exp(-0.5 * ((x - mean) / std) ** 2) / (std * math.sqrt(2.0 * math.pi)) + + def combine_classes(seg, classes): _seg = torch.zeros_like(seg) for i, c in enumerate(classes): diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index 131d908..c9cf7c6 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -1,6 +1,7 @@ +import functools import math from collections.abc import Callable, Sequence -from typing import Any, Protocol, Union, cast +from typing import Any, Protocol, Union import torch import torchvision.transforms._functional_tensor as F_t @@ -9,7 +10,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform -from smauglab.transforms.kernels import gaussian_kernel3d, laplace_kernel, scharr_kernels +from smauglab.transforms.kernels import depthwise_conv3d, gaussian_kernel3d, laplace_kernel, scharr_kernels, stacked_scharr_kernels from smauglab.transforms.rng import shared_choice @@ -126,7 +127,11 @@ def _select_and_check( if seg_mask is not None: region_mode = _choose_region_mode(transform.in_seg, transform.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=transform.mix_in_out) - if torch.isnan(x).any() or torch.isinf(x).any(): + # One `isfinite(...).all()` rather than `isnan().any() or isinf().any()`: each of + # those reads a reduction back to the host, and the `or` makes the second one + # unavoidable on the common (finite) path. This guard runs for every channel of + # every region-selecting transform. + if not bool(torch.isfinite(x).all()): print(f"Warning nan: {type(transform).__name__}{note}", flush=True) return None return x @@ -164,9 +169,15 @@ def _foreground(mask: torch.Tensor, dim: int) -> torch.Tensor: patch is not. So full coverage plus more than one channel means channel 0 is background, and the reduction skips it. """ - if mask.shape[dim] > 1 and bool((mask.amax(dim=dim) > 0).all()): - return mask.narrow(dim, 1, mask.shape[dim] - 1).amax(dim=dim) > 0 - return mask.amax(dim=dim) > 0 + covered = mask.amax(dim=dim) > 0 + if mask.shape[dim] == 1: + return covered + # `torch.where` rather than `if bool(...)`: reading that test back to Python blocks + # the host on everything queued behind it, and this runs once per region-selecting + # transform per batch. Both branches are a reduction over the class axis, so + # computing the unused one costs a fraction of what the stall does. + without_background = mask.narrow(dim, 1, mask.shape[dim] - 1).amax(dim=dim) > 0 + return torch.where(covered.all(), without_background, covered) def _apply_region_mode( @@ -198,13 +209,14 @@ def _apply_region_mode( transformed_max = torch.amax(transformed, dim=tuple(range(1, transformed.dim())), keepdim=True) transformed = (transformed - transformed_min) / (transformed_max - transformed_min + 1e-8) * (orig_max - orig_min) + orig_min - m = seg_mask.to(transformed.dtype).clone() + # No `.clone()`: nothing below writes into `m`, and on a [B, C, D, H, W] mask + # that copy is the size of the batch. + m = seg_mask.to(transformed.dtype) if mix_in_out: - for i in range(seg_mask.shape[0]): - # Create a tensor with random one and zero - - o = torch.randint(0, 2, (seg_mask.shape[1],), device=seg_mask.device, dtype=seg_mask.dtype) - m[i] = m[i] * o.view(-1, 1, 1, 1) # Broadcasting o to match the dimensions of m + # One keep/drop draw per (sample, class), in a single call rather than one + # `randint` per sample: the same draws in the same order, one kernel. + keep = torch.randint(0, 2, seg_mask.shape[:2], device=seg_mask.device, dtype=m.dtype) + m = m * keep.view(*seg_mask.shape[:2], *([1] * (m.dim() - 2))) m = _foreground(m, dim=1) m = m.to(transformed.dtype) @@ -219,10 +231,10 @@ def _apply_region_mode( transformed_max = torch.amax(transformed) transformed = (transformed - transformed_min) / (transformed_max - transformed_min + 1e-8) * (orig_max - orig_min) + orig_min - m = seg_mask.to(transformed.dtype).clone() + m = seg_mask.to(transformed.dtype) if mix_in_out: # Create a tensor with random one and zero - o = torch.randint(0, 2, (seg_mask.shape[0],), device=seg_mask.device, dtype=seg_mask.dtype) + o = torch.randint(0, 2, (seg_mask.shape[0],), device=seg_mask.device, dtype=m.dtype) m = m * o.view(-1, 1, 1, 1) # Broadcasting o to match the dimensions of m m = _foreground(m, dim=0) m = m.to(transformed.dtype) @@ -329,6 +341,15 @@ def get_kernel(self, device: torch.device) -> Union[Tensor, list[Tensor]]: raise NotImplementedError("Kernel type not implemented.") return kernel + def _per_sample_kernels(self, batch_size: int, device: torch.device) -> list[Tensor]: + """One freshly drawn kernel per sample, for the `same_on_batch=False` paths.""" + kernels = [] + for _ in range(batch_size): + drawn = self.get_kernel(device=device) + assert isinstance(drawn, Tensor) + kernels.append(drawn) + return kernels + @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Initialize kernel @@ -364,13 +385,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ assert isinstance(kernel, Tensor) x = apply_convolution(channel_data, kernel, dim=3) else: - x = torch.stack( - [ - apply_convolution(channel_data[b : b + 1], cast(Tensor, self.get_kernel(device=input.device)), dim=3).squeeze(0) - for b in range(channel_data.shape[0]) - ], - dim=0, - ) + x = apply_convolution_per_sample(channel_data, self._per_sample_kernels(channel_data.shape[0], input.device)) elif self.kernel_type == "UnsharpMask": # blur selected channel, compute mask and add scaled mask back Isharp​=I+α(I−G​∗I) if self.same_on_batch: @@ -378,25 +393,22 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ blurred = apply_convolution(channel_data, kernel, dim=3) unsharp_amount = torch.rand(1, device=input.device) * self.unsharp_amount else: - blurred = torch.stack( - [ - apply_convolution(channel_data[b : b + 1], cast(Tensor, self.get_kernel(device=input.device)), dim=3).squeeze(0) - for b in range(channel_data.shape[0]) - ], - dim=0, - ) + blurred = apply_convolution_per_sample(channel_data, self._per_sample_kernels(channel_data.shape[0], input.device)) amount_shape = [channel_data.shape[0]] + [1] * (channel_data.dim() - 1) unsharp_amount = torch.rand(channel_data.shape[0], device=input.device).view(amount_shape) * self.unsharp_amount mask = channel_data - blurred x = channel_data + unsharp_amount * mask elif self.kernel_type == "Scharr": - tot_ = torch.zeros_like(channel_data, device=input.device) - for k in kernel: - if self.absolute: - tot_ += torch.abs(apply_convolution(channel_data, k, dim=3)) - else: - tot_ += apply_convolution(channel_data, k, dim=3) - x = tot_ + # One convolution with three filters, not three convolutions summed: + # the three directional kernels differ only in their weights, so they + # fit in the output-channel axis of a single grouped conv. Same + # arithmetic, one pad and one kernel launch instead of three each. + weight = stacked_scharr_kernels(3, input.device, channel_data.dtype) + padded = F.pad(channel_data.unsqueeze(1), [1] * 6, mode="reflect") + directional = depthwise_conv3d(padded, weight, filters_per_plane=weight.shape[0]) + # `vector_norm(ord=1)` is abs-then-sum in one pass; `.abs().sum()` writes a + # second [N, 3, D, H, W] tensor out to memory and reads it straight back. + x = torch.linalg.vector_norm(directional, ord=1, dim=1) if self.absolute else directional.sum(dim=1) elif self.kernel_type == "RandConv": # One kernel for the batch under same_on_batch, a fresh one per # sample otherwise. This used to draw per sample unconditionally, @@ -406,16 +418,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ assert isinstance(kernel, Tensor) x = apply_convolution(channel_data, kernel, dim=3) else: - out = [] - for b in range(channel_data.shape[0]): - per_sample = self.get_kernel(device=input.device) - assert isinstance(per_sample, Tensor) - - conv = apply_convolution(channel_data[b : b + 1], per_sample, dim=3).squeeze(0) - - out.append(conv) - - x = torch.stack(out, dim=0) + x = apply_convolution_per_sample(channel_data, self._per_sample_kernels(channel_data.shape[0], input.device)) # Mix with original based on mix_prob, per sample. # @@ -674,12 +677,54 @@ def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torc if dim == 2: # noqa: SIM108 -- the 2d/3d split reads better spelled out than as a ternary img = F.conv2d(img, kernel, groups=img.shape[-(1 + dim)]) else: # dim == 3 - img = F.conv3d(img, kernel, groups=img.shape[-(1 + dim)]) + # Via `depthwise_conv3d` rather than `F.conv3d` directly. Every row of + # `kernel` is the same filter -- it came from an `expand` above -- so one row + # is handed over and the helper broadcasts it back across the planes, which + # is also what lets it make `groups > 1` out of a single-plane call. The + # bucketed pipelines hand every transform one sample at a time, and that is + # precisely the case cuDNN serves 35x slower than the depthwise kernel. + leading = img.shape[:-3] + planes = int(torch.tensor(leading).prod()) if leading else 1 + out = depthwise_conv3d(img.reshape(planes, 1, *img.shape[-3:]), kernel[:1, :1]) + img = out.reshape(*leading, *out.shape[-3:]) img = F_t._cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) return img +def apply_convolution_per_sample(channel_data: torch.Tensor, kernels: Sequence[Tensor]) -> torch.Tensor: + """Convolve `[N, D, H, W]` with one 3-D kernel per sample, in a single conv. + + The alternative -- and what every caller here used to do -- is a Python loop + calling `apply_convolution` on `channel_data[b : b + 1]`. That is a + single-channel `F.conv3d`, which cuDNN serves from its implicit-GEMM path: 4.3 ms + per sample for a 3x3x3 kernel over a 128^3 patch on an A40, against 0.12 ms for + the whole batch here. See `kernels.depthwise_conv3d` for why the shape matters. + + Kernels of different sizes are zero-padded up to the widest one. That is exact, + not an approximation: a zero tap contributes nothing, and `reflect` padding gives + the same value at a given offset however wide the pad is, so the extra taps read + real (if irrelevant) voxels and multiply them by zero. `RandConv` draws its kernel + size per sample, so without this the batch could not be convolved in one call. + """ + sizes = [int(k.shape[-1]) for k in kernels] + widest = max(sizes) + stacked = [] + for kernel, size in zip(kernels, sizes): + if size == widest: + stacked.append(kernel) + else: + lo = (widest - size) // 2 + padded = torch.zeros((widest, widest, widest), device=kernel.device, dtype=kernel.dtype) + padded[lo : lo + size, lo : lo + size, lo : lo + size] = kernel + stacked.append(padded) + weight = torch.stack(stacked).unsqueeze(1) # [N, 1, k, k, k] + + pad = widest // 2 + padded_img = F.pad(channel_data.unsqueeze(1), [pad] * 6, mode="reflect") + return depthwise_conv3d(padded_img, weight).squeeze(1) + + ## Noise transform @register( aug_id=AugId.GAUSSIAN_NOISE, @@ -740,8 +785,9 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ else: std = torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * self.std noise = torch.randn_like(input[:, c], device=input.device, dtype=input.dtype) - for i in range(input.shape[0]): - noise[i] = noise[i] * std[i] + self.mean + # Broadcast the per-sample std instead of writing one row at a time: + # same draws, same arithmetic, one kernel instead of N. + noise = noise * std.view(-1, *([1] * (noise.dim() - 1))) + self.mean orig = input[:, c] x = orig + noise @@ -819,9 +865,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ * (self.brightness_range[1] - self.brightness_range[0]) + self.brightness_range[0] ) - x = channel_data.clone() - for i in range(input.shape[0]): - x[i] = x[i] * factor[i] + x = channel_data * factor.view(-1, *([1] * (channel_data.dim() - 1))) checked = _select_and_check(self, orig, x, seg_mask) if checked is None: continue @@ -1076,19 +1120,15 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ torch.rand(1, device=input.device, dtype=input.dtype) * (self.contrast_range[1] - self.contrast_range[0]) + self.contrast_range[0] ) - x = channel_data.clone() - for i in range(input.shape[0]): - mean = x[i].mean() - x[i] = (x[i] - mean) * factor + mean + mean = channel_data.mean(dim=tuple(range(1, channel_data.dim())), keepdim=True) + x = (channel_data - mean) * factor + mean else: factor = ( torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * (self.contrast_range[1] - self.contrast_range[0]) + self.contrast_range[0] ) - x = channel_data.clone() - for i in range(input.shape[0]): - mean = x[i].mean() - x[i] = (x[i] - mean) * factor[i] + mean + mean = channel_data.mean(dim=tuple(range(1, channel_data.dim())), keepdim=True) + x = (channel_data - mean) * factor.view(-1, *([1] * (channel_data.dim() - 1))) + mean if self.retain_stats: x = _restore_stats(x, stats) @@ -1340,21 +1380,29 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ input = input.clone() seg_mask = params.get("seg") for c in self.apply_to_channel: + # The per-sample reductions -- max, and the two retain_stats passes -- are + # the expensive part and they are independent across samples, so they run + # once for the whole channel. What stays in the loop is the mix draw and + # the region selection, which both have to keep their per-sample identity: + # `_choose_region_mode` picks `in` / `out` / `all` for each sample + # separately, and hoisting it would silently make that one decision for the + # batch. Keeping the draws in the loop also keeps the RNG stream unchanged. + channel = input[:, c] + reduce_dims = tuple(range(1, channel.dim())) + keep_shape = (-1, *([1] * (channel.dim() - 1))) + inverted = channel.amax(dim=reduce_dims, keepdim=True) - channel + + if self.retain_stats: + eps = 1e-8 + orig_means = channel.mean(dim=reduce_dims).view(keep_shape) + orig_stds = channel.std(dim=reduce_dims).view(keep_shape) + new_mean = inverted.mean(dim=reduce_dims).view(keep_shape) + new_std = inverted.std(dim=reduce_dims).view(keep_shape) + inverted = (inverted - new_mean) / (new_std + eps) * orig_stds + orig_means + for i in range(input.shape[0]): - x = input[i, c] # shape [...spatial...] - orig = x.clone() - if self.retain_stats: - orig_means = x.mean() - orig_stds = x.std() - max_val = x.max() - x = max_val - x - - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - new_mean = x.mean() # scalar - new_std = x.std() # scalar - x = (x - new_mean) / (new_std + eps) * orig_stds + orig_means + orig = channel[i] + x = inverted[i] # Mix with original based on mix_prob if torch.rand(1).item() < self.mix_prob: @@ -1370,6 +1418,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ ## Histogram transform +#: Bin count for the histogram-equalisation CDF. Was a literal, repeated four times +#: across the index arithmetic and the clamp that has to agree with it. +_HIST_BINS = 256 + + @register( aug_id=AugId.HISTOGRAM_EQUAL, backend=Backend.GPU, @@ -1433,33 +1486,45 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ if self.retain_stats: stats = _channel_stats(channel_data) - # Process each batch element independently + # Equalise the whole batch at once. + # + # This was a per-sample loop around `torch.histc`, which needs its range as + # Python floats -- so every sample paid two `.item()` calls, and each of + # those blocks the host until the queue drains. The histogram is built here + # by scattering the bin indices the lookup below already needs, which is the + # same binning `histc(bins=256, min, max)` performs and removes the stalls + # along with the loop. batch_size = channel_data.shape[0] + flat = channel_data.reshape(batch_size, -1).to(torch.float32) + img_min = flat.amin(dim=1, keepdim=True) + img_max = flat.amax(dim=1, keepdim=True) + + bin_width = (img_max - img_min) / _HIST_BINS + indices = ((flat - img_min) / (bin_width + 1e-10)).long().clamp_(0, _HIST_BINS - 1) + + # `bincount` over row-offset indices, not `scatter_add_`: with only 256 bins + # a scatter has every thread in the block contending for the same handful + # of addresses, and measured 0.97 ms against bincount's 0.36 ms for a + # 2 x 128^3 batch. Counts are integers, so neither is approximate. + offset = indices + torch.arange(batch_size, device=indices.device).view(-1, 1) * _HIST_BINS + hist = torch.bincount(offset.reshape(-1), minlength=batch_size * _HIST_BINS).view(batch_size, _HIST_BINS) + + cdf = hist.cumsum(dim=1).to(flat.dtype) + # The smallest non-zero entry per sample. `cdf[cdf > 0].min()` cannot be + # written per row, so the empty bins are masked to +inf and reduced; a + # sample with no non-zero bin at all (an empty volume) falls back to the + # plain minimum, as the scalar version did. + positive = torch.where(cdf > 0, cdf, torch.full_like(cdf, float("inf"))) + cdf_min = positive.amin(dim=1, keepdim=True) + cdf_min = torch.where(torch.isinf(cdf_min), cdf.amin(dim=1, keepdim=True), cdf_min) + cdf = (cdf - cdf_min) / (cdf[:, -1:] - cdf_min + 1e-10) # Normalize to [0,1] + cdf = cdf * (img_max - img_min) + img_min # Scale back to image range + + channel_data = cdf.gather(1, indices).reshape(channel_data.shape).to(channel_data.dtype) + + # Mix with original based on mix_prob, per sample. The draws stay in a loop + # so the RNG stream is the one the scalar version produced. for b in range(batch_size): - img_b = channel_data[b] # Single image from batch [...spatial...] - - img_min, img_max = img_b.min(), img_b.max() - - # Flatten the image and compute the histogram - img_flattened = img_b.flatten().to(torch.float32) - hist = torch.histc(img_flattened, bins=256, min=img_min.item(), max=img_max.item()) - - # Compute the normalized cumulative distribution function (CDF) - cdf = hist.cumsum(dim=0) - cdf_min = cdf[cdf > 0].min() if (cdf > 0).any() else cdf.min() - cdf = (cdf - cdf_min) / (cdf[-1] - cdf_min + 1e-10) # Normalize to [0,1] - cdf = cdf * (img_max - img_min) + img_min # Scale back to image range - - # Compute bin edges and indices - bin_width = (img_max - img_min) / 256 - indices = ((img_flattened - img_min) / (bin_width + 1e-10)).long() - indices = torch.clamp(indices, 0, 255) - - # Perform histogram equalization - img_eq = cdf[indices] - channel_data[b] = img_eq.reshape(img_b.shape) - - # Mix with original based on mix_prob if torch.rand(1).item() < self.mix_prob: alpha = torch.rand(1, device=input.device) channel_data[b] = alpha * orig[b] + (1 - alpha) * channel_data[b] @@ -1475,6 +1540,48 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ return input +@functools.cache +def _polynomial_slot_map(order: int, dim: int, device: torch.device) -> Tensor: + """Flat coefficient index for every monomial, or -1 where the order forbids it. + + Shaped `(order+1,) * dim`, indexed by the per-axis exponents in the same + `x, y[, z]` order `RandomBiasFieldGPU._num_coeffs` counts them. + """ + size = order + 1 + slots = torch.full((size,) * dim, -1, dtype=torch.long) + index = 0 + if dim == 3: + for xo in range(size): + for yo in range(size - xo): + for zo in range(size - (xo + yo)): + slots[xo, yo, zo] = index + index += 1 + elif dim == 2: + for xo in range(size): + for yo in range(size - xo): + slots[xo, yo] = index + index += 1 + else: + raise ValueError("Only 2D or 3D spatial dims supported for bias field") + return slots.to(device) + + +@functools.cache +def _axis_power_tables(order: int, spatial: tuple[int, ...], device: torch.device, dtype: torch.dtype) -> tuple[Tensor, ...]: + """`[order+1, axis_length]` tables of `coordinate ** exponent`, one per axis. + + Returned innermost-axis-first (x, y[, z]) to match the exponent order of + `_polynomial_slot_map`. Cached: the tables depend only on the patch shape, which + is fixed for a training run, while the coefficients are redrawn every call. + """ + exponents = torch.arange(order + 1, device=device, dtype=dtype).view(-1, 1) + axes = [] + for length in reversed(spatial): # spatial is (D, H, W); x runs along W + coordinate = torch.linspace(-1, 1, length, device=device, dtype=dtype).view(1, -1) + axes.append(coordinate.pow(exponents)) + return tuple(axes) + + @register( aug_id=AugId.BIAS_FIELD, backend=Backend.GPU, @@ -1577,6 +1684,20 @@ def _make_grids(self, spatial_shape: tuple[int, ...], device: torch.device, dtyp else: raise ValueError("Spatial dims must be 2 or 3 for bias field") + def _coefficient_cube(self, coeffs: torch.Tensor, dim: int) -> torch.Tensor: + """Scatter the flat `(n_coeffs, B)` draw into a dense `(B, order+1, ...)` cube. + + The flat order is the nested `xo / yo / zo` walk the term loop used, so the + coefficient a given monomial gets is unchanged; the cube is zero wherever + `xo + yo + zo > order`, which is exactly the monomials that walk skipped. + """ + order = self.order + slot = _polynomial_slot_map(order, dim, coeffs.device) + used = slot >= 0 + cube = torch.zeros((coeffs.shape[1], *slot.shape), device=coeffs.device, dtype=coeffs.dtype) + cube[:, used] = coeffs[slot[used]].transpose(0, 1) + return cube + @torch.no_grad() def apply_transform( self, @@ -1601,35 +1722,27 @@ def apply_transform( dtype = input.dtype coeffs = self._sample_coeffs(batch_size, device, dtype, dim) # (n_coeffs, B) - grids = self._make_grids(spatial, device, dtype) seg_mask = params.get("seg") - # Initialize bias map per batch element - bias_map = torch.zeros((batch_size, *spatial), device=device, dtype=dtype) - - idx = 0 + # Evaluate the polynomial as a separable contraction rather than a term loop. + # + # Every monomial is a product of three one-dimensional powers, so the sum over + # monomials factorises: contract the coefficient cube against the axis power + # tables one axis at a time and the full-volume work collapses to the last + # step, a single matrix multiply. The loop this replaces materialised one + # [B, D, H, W] temporary per monomial -- twenty of them at the default order 3, + # each read and written in full -- to accumulate the same sum. + cube = self._coefficient_cube(coeffs, dim) + powers = _axis_power_tables(self.order, spatial, device, dtype) if dim == 3: - xg, yg, zg = grids # each shape (D,H,W) - for xo in range(self.order + 1): - x_term = xg.pow(xo) if xo > 0 else 1.0 - for yo in range(self.order + 1 - xo): - y_term = yg.pow(yo) if yo > 0 else 1.0 - for zo in range(self.order + 1 - (xo + yo)): - z_term = zg.pow(zo) if zo > 0 else 1.0 - # term shape (D,H,W) - term = x_term * y_term * z_term - # Add coefficient * term for each batch element - bias_map += coeffs[idx].view(-1, *([1] * dim)) * term # broadcast over spatial - idx += 1 + x_pow, y_pow, z_pow = powers # each (order+1, axis_length) + partial = torch.einsum("bxyz,zd->bxyd", cube, z_pow) + partial = torch.einsum("bxyd,yh->bxdh", partial, y_pow) + bias_map = torch.einsum("bxdh,xw->bdhw", partial, x_pow) else: # dim == 2 - xg, yg = grids # (H,W) - for xo in range(self.order + 1): - x_term = xg.pow(xo) if xo > 0 else 1.0 - for yo in range(self.order + 1 - xo): - y_term = yg.pow(yo) if yo > 0 else 1.0 - term = x_term * y_term # (H,W) - bias_map += coeffs[idx].view(-1, *([1] * dim)) * term - idx += 1 + x_pow, y_pow = powers + partial = torch.einsum("bxy,yh->bxh", cube, y_pow) + bias_map = torch.einsum("bxh,xw->bhw", partial, x_pow) # Exponential to ensure positive field bias_field = torch.exp(bias_map) # (N, *spatial) @@ -1665,6 +1778,38 @@ def apply_transform( return input +def _batched_quantiles(channel_data: torch.Tensor, lower: torch.Tensor, upper: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-sample `torch.quantile` at a per-sample probability, in one sort. + + `torch.quantile` applies every `q` it is given to every row, so a per-sample + probability forced a Python loop with two calls per sample -- and each call sorts + the row again. Here the rows are sorted once and both quantiles are read off by + index, with `quantile`'s own linear interpolation + (`lower + frac * (upper - lower)`), so the values are the ones it would return. + + The default `max_clamp_amount=0` asks for the 0th and 100th percentile, which is + the volume's own range and makes the clamp a no-op; that case skips the sort. + """ + batch_size = channel_data.shape[0] + flat = channel_data.reshape(batch_size, -1) + if bool((lower == 0).all()) and bool((upper == 1).all()): + return flat.amin(dim=1), flat.amax(dim=1) + + ordered, _ = torch.sort(flat, dim=1) + last = ordered.shape[1] - 1 + + def pick(prob: torch.Tensor) -> torch.Tensor: + position = prob.clamp(0.0, 1.0).to(ordered.dtype) * last + low_idx = position.floor().long().clamp(0, last) + high_idx = position.ceil().long().clamp(0, last) + frac = (position - low_idx.to(position.dtype)).unsqueeze(1) + low_val = ordered.gather(1, low_idx.unsqueeze(1)) + high_val = ordered.gather(1, high_idx.unsqueeze(1)) + return (low_val + frac * (high_val - low_val)).squeeze(1) + + return pick(lower), pick(upper) + + # Random clamping transform @register( aug_id=AugId.CLAMP, @@ -1725,22 +1870,20 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ if self.retain_stats: stats = _channel_stats(channel_data) + batch_size = input.shape[0] if self.same_on_batch: - min_percentile = torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount - max_percentile = 1.0 - (torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount) - x = channel_data.clone() - for i in range(input.shape[0]): - min_val = torch.quantile(x[i].flatten(), min_percentile) - max_val = torch.quantile(x[i].flatten(), max_percentile) - x[i] = torch.clamp(x[i], min_val, max_val) + min_percentile = (torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount).expand(batch_size) + max_percentile = (1.0 - (torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount)).expand(batch_size) else: - x = channel_data.clone() - for i in range(input.shape[0]): - min_percentile = torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount - max_percentile = 1.0 - (torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount) - min_val = torch.quantile(x[i].flatten(), min_percentile) - max_val = torch.quantile(x[i].flatten(), max_percentile) - x[i] = torch.clamp(x[i], min_val, max_val) + # [N, 2] rather than 2N scalar draws: the row-major fill order is the + # same (min, max) per sample the loop produced. + draws = torch.rand(batch_size, 2, device=input.device, dtype=input.dtype) * self.max_clamp_amount + min_percentile = draws[:, 0] + max_percentile = 1.0 - draws[:, 1] + + lo, hi = _batched_quantiles(channel_data, min_percentile, max_percentile) + bounds_shape = (-1, *([1] * (channel_data.dim() - 1))) + x = torch.clamp(channel_data, lo.view(bounds_shape), hi.view(bounds_shape)) if self.retain_stats: x = _restore_stats(x, stats) diff --git a/smauglab/transforms/gpu/domain_transfer.py b/smauglab/transforms/gpu/domain_transfer.py index 9c036bf..33d9fec 100644 --- a/smauglab/transforms/gpu/domain_transfer.py +++ b/smauglab/transforms/gpu/domain_transfer.py @@ -251,7 +251,14 @@ def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> tuple[ @staticmethod def _accumulate(lut_used: Tensor, per_class: bool, w_b: Tensor, il: Tensor, ih: Tensor, xf: Tensor, n_seg_c: int) -> Tensor: - """Class-weighted LUT interpolation ``Σ_c w_c · interp(LUT_c, x)`` → ``[D, H, W]``.""" + """Class-weighted LUT interpolation ``Σ_c w_c · interp(LUT_c, x)`` → ``[D, H, W]``. + + Left as a loop over classes on purpose. Reading every class at once -- + `lut_all[:, il]` for a `[C, D, H, W]` gather, then one fused reduction -- + moves the same number of bytes and was measured 3% *slower* on the shipped + four-class bank, because the fused expression needs `[C, D, H, W]` + temporaries where the loop needs single-volume ones. + """ acc = torch.zeros_like(xf) for c in range(n_seg_c): lut_c = lut_used[c if per_class else 0, c] # [L] @@ -263,7 +270,10 @@ def _is_zscore(self, x: Tensor) -> bool: return True if self.zscore_io == "never": return False - return bool((x.min() < -1e-3) or (x.max() > 1.0 + 1e-3)) + # One reduction and one host read: `x.min()` and `x.max()` were separate + # passes over the patch, and `or` on two device scalars synchronises twice. + low, high = torch.aminmax(x) + return bool((low < -1e-3) | (high > 1.0 + 1e-3)) def _to_unit(self, x: Tensor) -> Tensor: """Map a (z-scored) image into the LUT's [0,1] domain via percentile scaling (clip), @@ -271,8 +281,10 @@ def _to_unit(self, x: Tensor) -> Tensor: flat = x.reshape(-1).float() if flat.numel() > 1_000_000: # cap for torch.quantile flat = flat[torch.linspace(0, flat.numel() - 1, 1_000_000, device=x.device).long()] - lo = torch.quantile(flat, self.pct / 100.0) - hi = torch.quantile(flat, 1.0 - self.pct / 100.0) + # Both cut points from one call, so the million values are sorted once + # instead of twice. + probabilities = torch.tensor([self.pct / 100.0, 1.0 - self.pct / 100.0], device=x.device, dtype=flat.dtype) + lo, hi = torch.quantile(flat, probabilities).unbind() return ((x - lo) / (hi - lo).clamp_min(1e-6)).clamp(0.0, 1.0) @torch.no_grad() diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index dea6d87..225623d 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -1,3 +1,5 @@ +import functools +import math from collections.abc import Sequence from typing import Any @@ -14,17 +16,58 @@ def _kmeans_1d(values: torch.Tensor, C: int, n_iter: int = 10) -> torch.Tensor: - """1-D K-means on foreground values. Returns (C,) centroids.""" - centroids = torch.linspace(values.min().item(), values.max().item(), C, device=values.device) + """1-D K-means on foreground values. Returns (C,) centroids. + + Lloyd's algorithm on sorted data, which is the same algorithm the dense version + ran but with each step costing O(C) instead of O(n*C): + + * In one dimension with sorted centroids, nearest-centroid assignment is an + interval split at the midpoints, so the cluster boundaries come from a + `searchsorted` over the sorted values rather than an `argmin` over an + `n x C` distance matrix. + * Cluster sums then come from differences of one prefix sum, instead of a + `scatter_add_` of n values into C bins -- which on CUDA means every thread + contending for one of a handful of addresses, and measured as the single most + expensive operation in `RandomPaletteGPU`. + * The `torch.allclose` convergence test is gone. It was a host synchronisation + on every iteration, and the loop body it guards is now a few kernels over C + elements; waiting for the GPU cost more than the iterations it saved. + + The prefix sum is accumulated in float64. A float32 running total over ten + thousand values loses enough low-order bits that two nearby cluster boundaries + can disagree about which side a value fell on, and the result is an augmentation + whose region map depends on the subsample size. + """ + device = values.device + v_min, v_max = values.min(), values.max() + if C <= 1: + return v_min.reshape(1) + + # `torch.linspace(a, b, C)`, written out: the last point is `b` exactly, not + # `a + (C-1)*step`, and starting K-means from a different point can land it in a + # different local optimum. + step = (v_max - v_min) / (C - 1) + centroids = v_min + step * torch.arange(C, device=device, dtype=values.dtype) + centroids = torch.cat([centroids[:-1], v_max.reshape(1)]) + + ordered, _ = torch.sort(values) + prefix = torch.cat([torch.zeros(1, device=device, dtype=torch.float64), ordered.double().cumsum(0)]) + n = ordered.numel() + tail = torch.full((1,), n, dtype=torch.long, device=device) + head = torch.zeros(1, dtype=torch.long, device=device) + for _ in range(n_iter): - d = torch.abs(values.unsqueeze(1) - centroids.unsqueeze(0)) - lbl = torch.argmin(d, dim=1) - s = torch.zeros(C, device=values.device).scatter_add_(0, lbl, values) - n = torch.zeros(C, device=values.device).scatter_add_(0, lbl, torch.ones_like(values)) - new_c = torch.where(n > 0, s / n, centroids) - if torch.allclose(centroids, new_c): - break - centroids = new_c + boundaries = (centroids[:-1] + centroids[1:]) / 2.0 + # `right=True`: a value exactly on a midpoint joins the lower cluster, which + # is what `argmin` did (it returns the first of two equal distances) and what + # the `torch.bucketize` the caller runs afterwards does. + edges = torch.searchsorted(ordered, boundaries, right=True) + starts = torch.cat([head, edges]) + ends = torch.cat([edges, tail]) + counts = ends - starts + sums = prefix[ends] - prefix[starts] + means = (sums / counts.clamp_min(1)).to(centroids.dtype) + centroids = torch.where(counts > 0, means, centroids) return centroids @@ -54,27 +97,45 @@ def _voronoi_region_ids( """ N = lbl_l.shape[0] rid = torch.zeros(N, dtype=torch.long, device=device) + + # Everything the loop needs off the GPU, in two reads instead of three per + # cluster. Each `.item()` inside the loop drained the queue, and with up to six + # clusters per sample and two samples per batch that was the dominant cost of + # the whole transform on a 128^3 patch. + # + # Drawing the decisions up front consumes a fixed two uniforms per cluster where + # the scalar version drew the second one only when the first said "subdivide", + # so a seeded run no longer reproduces the previous release's region maps. The + # draws are the same independent uniforms either way. + counts = torch.bincount(lbl_l[fg > 0], minlength=C)[:C].tolist() + decisions = torch.rand(C, 2, device=device).tolist() + offset = 0 for c in range(C): - c_mask = lbl_l == c - c_fg_mask = c_mask & (fg > 0) - n_fg = int(c_fg_mask.sum().item()) + n_fg = int(counts[c]) if n_fg == 0: continue - if n_fg < 2 or torch.rand(1, device=device).item() < skip_sub_parc_prob: + c_mask = lbl_l == c + if n_fg < 2 or decisions[c][0] < skip_sub_parc_prob: S = 1 else: - s_idx = int(torch.rand(1, device=device).item() * len(s_choices)) + s_idx = int(decisions[c][1] * len(s_choices)) S = min(s_choices[s_idx], n_fg) if S <= 1: - rid = torch.where(c_mask, torch.full_like(rid, offset), rid) + rid.masked_fill_(c_mask, offset) offset += 1 continue - seed_idx = torch.multinomial(c_fg_mask.float(), S, replacement=False) - centroids_v = coords.index_select(0, seed_idx) - d = torch.cdist(coords, centroids_v) - sub = torch.argmin(d, dim=1) - rid = torch.where(c_mask, offset + sub, rid) + # Only this cluster's voxels take part: the result is written under `c_mask` + # anyway, so the full-volume `cdist` computed C times as many distances as + # the output has rows. + member = c_mask.nonzero(as_tuple=True)[0] + fg_member = member[fg[member] > 0] + # The S largest of n iid uniforms is a uniformly random S-subset, and `topk` + # for S <= 10 is a reduction where `randperm` is a full sort of n: 0.11 ms + # against 0.49 ms at a million foreground voxels. + seed_idx = fg_member[torch.topk(torch.rand(fg_member.numel(), device=device), S).indices] + d = torch.cdist(coords.index_select(0, member), coords.index_select(0, seed_idx)) + rid[member] = offset + torch.argmin(d, dim=1) offset += S # At least one region, always. Every cluster can be skipped above -- a constant image # has no foreground at all, so `n_fg` is zero for all of them -- and `offset` then @@ -88,11 +149,82 @@ def _voronoi_region_ids( # ───────────────────────────────────────────────────────────────────────────── +def segment_sum(ids: torch.Tensor, values: torch.Tensor, n_segments: int) -> torch.Tensor: + """Sum `values` into `n_segments` bins indexed by `ids`. Returns `(n_segments,)`. + + `torch.zeros(R).scatter_add_(0, ids, values)` is the obvious spelling and the + slow one: with a couple of dozen bins and two million values, every thread in + the grid contends for one of a handful of global addresses. `bincount` keeps its + accumulator in shared memory per block, and measured 0.16 ms against 0.52 ms for + a 128^3 volume over 24 regions. + + `bincount` sizes its output from the largest id present, so a trailing segment + that no voxel landed in would come back missing; the pad restores it as a zero. + """ + summed = torch.bincount(ids, weights=values, minlength=n_segments).to(values.dtype) + if summed.numel() < n_segments: + summed = F.pad(summed, (0, n_segments - summed.numel())) + return summed[:n_segments] + + +def foreground_classes(labels: torch.Tensor) -> torch.Tensor: + """The distinct positive values in an integer label volume, ascending. + + `labels.unique()` radix-sorts the whole volume to answer a question about a + handful of small integers; counting them and keeping the non-empty bins is half + the time on a 2 x 128^3 batch, and the result is identical. + """ + flat = labels.reshape(-1) + counts = torch.bincount(flat.clamp_min(0)) + present = counts.nonzero().flatten() + return present[present > 0] + + +@functools.lru_cache(maxsize=8) +def voxel_coordinates(shape: tuple[int, int, int], device: torch.device) -> torch.Tensor: + """`[D*H*W, 3]` ijk voxel coordinates for a patch shape, built once and shared. + + This is a function of the patch shape alone -- fixed for a training run -- but + was rebuilt on every call, and at 128^3 the `meshgrid`/`stack` writes 24 MB each + time. The returned tensor is shared; treat it as read-only. + """ + depth, height, width = shape + return torch.stack( + torch.meshgrid( + torch.arange(depth, device=device, dtype=torch.float32), + torch.arange(height, device=device, dtype=torch.float32), + torch.arange(width, device=device, dtype=torch.float32), + indexing="ij", + ), + dim=-1, + ).reshape(depth * height * width, 3) + + +#: sqrt(2*pi), as a Python float. It used to be built with `torch.sqrt(torch.tensor( +#: 2*pi, device=...))` inside `_normal_pdf`, which is a host-to-device copy and a +#: kernel launch for a compile-time constant -- once per region per sample. +_SQRT_TWO_PI = math.sqrt(2.0 * math.pi) + + def _normal_pdf(x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: inv = 1.0 / (std + 1e-6) - return (inv / (torch.sqrt(torch.tensor(2.0 * 3.141592653589793, device=x.device, dtype=x.dtype)))) * torch.exp( - -0.5 * ((x - mean) * inv) ** 2 - ) + return (inv / _SQRT_TWO_PI) * torch.exp(-0.5 * ((x - mean) * inv) ** 2) + + +def _region_moments(region_mask: torch.Tensor, design: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Voxel count, mean and standard deviation per region, in one matrix multiply. + + `design` is the `[S, 3]` stack of `1`, `x` and `x**2`, so the product holds the + three sums each region's statistics need. The variance comes from + `E[x^2] - E[x]^2`, which is the same quantity the explicit + `((x - mean) * mask) ** 2` pass computed -- the mask is binary, so squaring it + changes nothing -- without materialising the centred array. + """ + sums = region_mask.to(design.dtype) @ design # (R, 3) + counts = sums[:, 0].clamp_min(1) + means = sums[:, 1] / counts + variances = (sums[:, 2] / counts - means * means).clamp_min(0) + return counts, means, variances.sqrt() ## Redistribute segmentation values transform (GPU) @@ -194,6 +326,19 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ denom = (img_max - img_min).clamp_min(1e-6) x_batch = (img_batch - img_min) / denom + # Everything the per-sample loop needs to read back to Python, read once. + # These were three `.item()` calls inside the loop -- the empty-mask test, + # the redistribution mode and the dilation count -- and each one stalls the + # host until the queue drains. Drawing the two random ones for the whole + # batch up front gives each sample its own independent draw exactly as + # before, but no longer in the interleaved order the loop produced, so a + # seeded run differs from the previous release's. + has_foreground = (seg.reshape(N, -1) > 0).any(dim=1).tolist() + in_seg_draws = (torch.rand(N, device=input.device) <= self.in_seg).tolist() + dilation_draws = torch.randint( + self.dilation_iterations_range[0], self.dilation_iterations_range[1] + 1, (N,), device=input.device + ).tolist() + # Iterate per sample (seg can differ in shape or labels per sample) for b in range(N): x = x_batch[b] @@ -206,13 +351,12 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # and a few cases are effectively unlabelled, so this fires on real # anatomy and produces a perfectly finite, badly out-of-distribution # patch that no NaN guard can see. - if (seg[b] > 0).sum() == 0: + if not has_foreground[b]: input[b, c] = img_batch[b] continue # Decide redistribution mode once per sample - # Scalar random flag for redistribution mode - in_seg_bool = torch.rand((), device=input.device) <= self.in_seg + in_seg_bool = in_seg_draws[b] # Binary masks for regions masks = regions[b] # (R, ...) @@ -220,12 +364,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # Vectorized dilation for all regions (3 iterations) dilated = masks.float() - dilation_iterations = int( - torch.randint(self.dilation_iterations_range[0], self.dilation_iterations_range[1] + 1, (1,), device=input.device)[ - 0 - ].item() - ) - for _ in range(dilation_iterations): + for _ in range(int(dilation_draws[b])): if spatial_dims == 3: dilated = F.max_pool3d(dilated.unsqueeze(0), 3, 1, 1).squeeze(0) else: @@ -237,21 +376,13 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ mask_flat = masks.view(R, -1) dil_flat = dilated_excl.view(R, -1) - # Region counts - counts = mask_flat.sum(dim=1).clamp_min(1) - # Means - means = (mask_flat * x_flat).sum(dim=1) / counts - # Std (compute variance then sqrt) avoid indexing overhead - diffs = (x_flat - means.view(R, 1)) * mask_flat - vars = (diffs * diffs).sum(dim=1) / counts.clamp_min(1) - stds = vars.sqrt() - - # Dilated stats - dil_counts = dil_flat.sum(dim=1).clamp_min(1) - dil_means = (dil_flat * x_flat).sum(dim=1) / dil_counts - dil_diffs = (x_flat - dil_means.view(R, 1)) * dil_flat - dil_vars = (dil_diffs * dil_diffs).sum(dim=1) / dil_counts - dil_stds = dil_vars.sqrt() + # Count, mean and variance per region, as one [R, S] x [S, 3] matrix + # multiply each. Spelled out, that was six passes over an [R, S] + # array and two [R, S] temporaries -- 32 MB apiece for four regions + # of a 128^3 patch -- to produce six numbers per region. + design = torch.stack([torch.ones_like(x_flat[0]), x_flat[0], x_flat[0] * x_flat[0]], dim=1) # (S, 3) + counts, means, stds = _region_moments(mask_flat, design) + dil_counts, dil_means, dil_stds = _region_moments(dil_flat, design) # redist_std per region std_noise_range = ( @@ -262,31 +393,30 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ torch.full((R,), 0.01, device=input.device, dtype=input.dtype), ) - # Build additive term + # Build additive term. + # + # Accumulated in place rather than stacked: the global branch used to + # build a list of R full-volume tensors and `torch.stack` them before + # summing, which is an extra R-volume allocation and an extra pass over + # it. The in-region branch indexed each region out with a boolean mask, + # and a boolean index is a `nonzero` -- a host synchronisation per + # region, on top of the gather. Multiplying by the mask writes the same + # values at the same voxels, because the regions are summed either way. + # + # The `counts[r] == 0` guards both branches carried could not fire: + # `counts` is `clamp_min(1)`. to_add = torch.zeros_like(x) rand_sign = 2 * torch.rand(R, device=input.device) - 1 # random sign factor per region - if in_seg_bool.item(): - # Only inside region - for r in range(R): - if counts[r] == 0: # skip empty - continue - region_vals = x[mask_flat[r].view(x.shape)] - pdf_vals = _normal_pdf(region_vals, means[r], redist_std[r]) * rand_sign[r] - to_add[mask_flat[r].view(x.shape)] += pdf_vals - else: - # Global additive influence per region - pdf_all = [] - for r in range(R): - if counts[r] == 0: - continue - pdf_all.append(_normal_pdf(x, means[r], redist_std[r]) * rand_sign[r]) - if pdf_all: - to_add += torch.stack(pdf_all, dim=0).sum(dim=0) - - # Normalize to_add if non-zero + for r in range(R): + contribution = _normal_pdf(x, means[r], redist_std[r]) * rand_sign[r] + if in_seg_bool: + contribution = contribution * mask_flat[r].view(x.shape) + to_add += contribution + + # Normalize to_add if non-zero. `torch.where`, not `if`: the test is one + # more host read, and both sides are cheap. tmin, tmax = to_add.min(), to_add.max() - if (tmax - tmin) > 1e-8: - x = x + 2 * (to_add - tmin) / (tmax - tmin + 1e-6) + x = torch.where(tmax - tmin > 1e-8, x + 2 * (to_add - tmin) / (tmax - tmin + 1e-6), x) # Restore stats if self.retain_stats: @@ -297,7 +427,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ x = x * orig_std[b] + orig_mean[b] # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): + if not bool(torch.isfinite(x).all()): print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[b, c] = x @@ -415,16 +545,9 @@ def apply_transform( flat_m_all = (images_01 > self.dark_threshold).float() # foreground mask - # Voxel coordinates (shared — same spatial dims for every sample) - coords = torch.stack( - torch.meshgrid( - torch.arange(D, device=device, dtype=torch.float32), - torch.arange(H, device=device, dtype=torch.float32), - torch.arange(W, device=device, dtype=torch.float32), - indexing="ij", - ), - dim=-1, - ).reshape(N, 3) + # Voxel coordinates (shared — same spatial dims for every sample, and the + # same from one call to the next, so they are built once per patch shape). + coords = voxel_coordinates((D, H, W), device) # ── Step 1: PALETTE K-means + Voronoi per-region affine remap ────────── synth_list = [] @@ -465,8 +588,8 @@ def apply_transform( self.skip_sub_parc_prob, ) - s_c = torch.zeros(R, device=device).scatter_add_(0, rid, flat * flat_m) - n_c = torch.zeros(R, device=device).scatter_add_(0, rid, flat_m) + s_c = segment_sum(rid, flat * flat_m, R) + n_c = segment_sum(rid, flat_m, R) mean_c = s_c / n_c.clamp(min=eps) mu_c = torch.rand(R, device=device) @@ -495,8 +618,7 @@ def apply_transform( labels = F.interpolate(labels.float(), size=(D, H, W), mode="nearest-exact").long() lbl = labels[:, 0].reshape(B, N).clamp(min=0) - unique_classes = lbl.unique() - unique_classes = unique_classes[unique_classes > 0] + unique_classes = foreground_classes(lbl) if self.label_classes is not None: keep = torch.tensor(self.label_classes, device=device) unique_classes = unique_classes[torch.isin(unique_classes, keep)] diff --git a/smauglab/transforms/gpu/palette/base.py b/smauglab/transforms/gpu/palette/base.py index ab77667..636ff2a 100644 --- a/smauglab/transforms/gpu/palette/base.py +++ b/smauglab/transforms/gpu/palette/base.py @@ -19,6 +19,8 @@ import torch from torch import nn +from smauglab.transforms.gpu.fromSeg import segment_sum + @dataclass(slots=True) class BlockContext: @@ -71,8 +73,8 @@ def signed_alpha_affine_remap( device = image01.device alpha_lo, alpha_hi = alpha_magnitude_range - s_c = torch.zeros(n_regions, device=device).scatter_add_(0, region_ids, image01 * fg_mask) - n_c = torch.zeros(n_regions, device=device).scatter_add_(0, region_ids, fg_mask) + s_c = segment_sum(region_ids, image01 * fg_mask, n_regions) + n_c = segment_sum(region_ids, fg_mask, n_regions) mean_c = s_c / n_c.clamp(min=eps) mu_c = torch.rand(n_regions, device=device) diff --git a/smauglab/transforms/gpu/palette/overlay.py b/smauglab/transforms/gpu/palette/overlay.py index 6b7b115..af2619a 100644 --- a/smauglab/transforms/gpu/palette/overlay.py +++ b/smauglab/transforms/gpu/palette/overlay.py @@ -6,6 +6,8 @@ from torch import nn from torch.nn import functional as F +from smauglab.transforms.gpu.fromSeg import foreground_classes + #: A scalar blend, or a [lo, hi] range sampled per label per sample. BlendSpec = float | Sequence[float] @@ -71,8 +73,7 @@ def remap_labels( labels = F.interpolate(labels.float(), size=(depth, height, width), mode="nearest-exact").long() lbl = labels[:, 0].reshape(batch, n_voxels).clamp(min=0) - classes = lbl.unique() - classes = classes[classes > 0] + classes = foreground_classes(lbl) if self.label_classes is not None: keep = torch.tensor(self.label_classes, device=device) classes = classes[torch.isin(classes, keep)] diff --git a/smauglab/transforms/gpu/palette/partitioners.py b/smauglab/transforms/gpu/palette/partitioners.py index 851b454..3bf803e 100644 --- a/smauglab/transforms/gpu/palette/partitioners.py +++ b/smauglab/transforms/gpu/palette/partitioners.py @@ -192,5 +192,12 @@ def _densify_region_ids(region_ids: torch.Tensor) -> tuple[torch.Tensor, int]: is sparse. `signed_alpha_affine_remap` indexes per-region tensors of length R with these, so a sparse id would index out of bounds. """ - unique, inverse = torch.unique(region_ids, return_inverse=True) - return inverse.long(), int(unique.numel()) + # A lookup table off a bincount, not `torch.unique(return_inverse=True)`: the ids + # are small non-negative integers, so counting the bins and numbering the + # non-empty ones gives the same ascending remap without radix-sorting two + # million voxels. + counts = torch.bincount(region_ids) + present = counts.nonzero().flatten() + lookup = torch.zeros(counts.numel(), dtype=torch.long, device=region_ids.device) + lookup[present] = torch.arange(present.numel(), device=region_ids.device) + return lookup[region_ids], int(present.numel()) diff --git a/smauglab/transforms/gpu/palette/transform.py b/smauglab/transforms/gpu/palette/transform.py index 8dd6304..87ebe49 100644 --- a/smauglab/transforms/gpu/palette/transform.py +++ b/smauglab/transforms/gpu/palette/transform.py @@ -8,7 +8,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform -from smauglab.transforms.gpu.fromSeg import _gaussian_blur_3d, collapse_onehot_to_index +from smauglab.transforms.gpu.fromSeg import _gaussian_blur_3d, collapse_onehot_to_index, voxel_coordinates from smauglab.transforms.gpu.palette.base import ( BlockContext, InitialPartitioner, @@ -125,16 +125,9 @@ def apply_transform( flat_m_all = (images_01 > self.dark_threshold).float() # foreground mask - # Voxel coordinates (shared — same spatial dims for every sample) - coords = torch.stack( - torch.meshgrid( - torch.arange(D, device=device, dtype=torch.float32), - torch.arange(H, device=device, dtype=torch.float32), - torch.arange(W, device=device, dtype=torch.float32), - indexing="ij", - ), - dim=-1, - ).reshape(N, 3) + # Voxel coordinates (shared — same spatial dims for every sample, and the + # same from one call to the next, so they are built once per patch shape). + coords = voxel_coordinates((D, H, W), device) # ── Step 1: the block stack builds a partition; the remap is fixed ───── synth_list = [] diff --git a/smauglab/transforms/gpu/spatial.py b/smauglab/transforms/gpu/spatial.py index cf564b1..61fe7ff 100644 --- a/smauglab/transforms/gpu/spatial.py +++ b/smauglab/transforms/gpu/spatial.py @@ -270,35 +270,36 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ interp_down = resample interp_up = resample - # Process per-channel and per-batch element - out = input.clone() - + # `scales` lives on the device, so each `float(...)` below used to block the + # host until everything queued had finished -- three times per sample. One + # `.tolist()` reads the whole [B, 3] draw back in a single transfer. + scale_rows = scales.tolist() + align = False if "linear" in interp_down else None + + # Samples that drew the same target size resample together: one + # `F.interpolate` over a [k, C, D, H, W] slab instead of k of them. With + # `same_on_batch` the whole batch is a single group, and the loop that + # remains has one iteration. + groups: dict[tuple[int, int, int], list[int]] = {} for b in range(batch_size): - x = input[b] # [C, D, H, W] - - sx, sy, sz = scales[b] - # compute downsampled size - down_D = max(1, round(float(sz) * D)) - down_H = max(1, round(float(sy) * H)) - down_W = max(1, round(float(sx) * W)) - - # downsample - x_down = F.interpolate( - x.unsqueeze(0), - size=(down_D, down_H, down_W), - mode=interp_down, - align_corners=False if "linear" in interp_down else None, - ) - - # upsample back to original resolution (keep as 4D tensor [1,1,D,H,W]) - x_up = F.interpolate( - x_down, - size=(D, H, W), - mode=interp_up, - align_corners=False if "linear" in interp_up else None, - ).squeeze(0) # [C, D, H, W] - - out[b] = x_up + sx, sy, sz = scale_rows[b] + size = (max(1, round(sz * D)), max(1, round(sy * H)), max(1, round(sx * W))) + groups.setdefault(size, []).append(b) + + # `empty_like`, not `clone`: every row is written below. + out = torch.empty_like(input) + for size, members in groups.items(): + # A one-member group is a view, not a gather: with a small batch most + # groups are singletons, and `index_select`/`index_copy_` would copy the + # sample in and out again for nothing. + if len(members) == 1: + b = members[0] + chunk = F.interpolate(input[b : b + 1], size=size, mode=interp_down, align_corners=align) + out[b : b + 1] = F.interpolate(chunk, size=(D, H, W), mode=interp_up, align_corners=align) + continue + index = torch.as_tensor(members, device=input.device) + chunk = F.interpolate(input.index_select(0, index), size=size, mode=interp_down, align_corners=align) + out.index_copy_(0, index, F.interpolate(chunk, size=(D, H, W), mode=interp_up, align_corners=align)) return out @@ -436,45 +437,38 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ interp_down = resample interp_up = resample - # Process per-channel and per-batch element - out = input.clone() + # One `.tolist()` instead of three device reads per sample, and one + # `F.interpolate` per distinct target size instead of one per (sample, + # channel) -- see the sibling in RandomLowResTransformGPU. + scale_rows = scales.tolist() + align = False if "linear" in interp_down else None + channels = list(self.apply_to_channel) + + groups: dict[tuple[int, int, int], list[int]] = {} for b in range(batch_size): - # start from the original per-sample tensor so we only overwrite selected channels - canvas = input[b].clone() - - for c in self.apply_to_channel: - x = input[b, c] # [D, H, W] - - sx, sy, sz = scales[b] - # compute downsampled size - down_D = max(1, round(float(sz) * D)) - down_H = max(1, round(float(sy) * H)) - down_W = max(1, round(float(sx) * W)) - - # downsample - x_down = F.interpolate( - x.unsqueeze(0).unsqueeze(0), - size=(down_D, down_H, down_W), - mode=interp_down, - align_corners=False if "linear" in interp_down else None, - ) - - # upsample back to original resolution - x_up = ( - F.interpolate( - x_down, - size=(D, H, W), - mode=interp_up, - align_corners=False if "linear" in interp_up else None, - ) - .squeeze(0) - .squeeze(0) - ) # [D, H, W] - - # place patch back into the canvas for the correct channel only - canvas[c] = x_up - - out[b] = canvas + sx, sy, sz = scale_rows[b] + size = (max(1, round(sz * D)), max(1, round(sy * H)), max(1, round(sx * W))) + groups.setdefault(size, []).append(b) + + # Only the selected channels are rewritten, so the untouched ones have to be + # carried over -- hence a clone rather than an empty tensor. + out = input.clone() + for size, members in groups.items(): + # A one-member group goes through views: with a small batch most groups + # are singletons, and gathering the sample in and scattering it back out + # costs more than the resampling it was meant to share. + if len(members) == 1: + b = members[0] + for c in channels: + chunk = F.interpolate(input[b, c][None, None], size=size, mode=interp_down, align_corners=align) + out[b, c] = F.interpolate(chunk, size=(D, H, W), mode=interp_up, align_corners=align)[0, 0] + continue + index = torch.as_tensor(members, device=input.device) + chunk = input.index_select(0, index).index_select(1, torch.as_tensor(channels, device=input.device)) + chunk = F.interpolate(chunk, size=size, mode=interp_down, align_corners=align) + chunk = F.interpolate(chunk, size=(D, H, W), mode=interp_up, align_corners=align) + for position, c in enumerate(channels): + out[index, c] = chunk[:, position] return out @@ -530,16 +524,22 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # generator (including its "at least one axis" guarantee) was dead code. flips = params.get("flip") + # `bool(flips[b, axis])` read one element of a device tensor back to the host + # at a time: three synchronisations per sample, for a transform whose actual + # work is a strided copy. One `.tolist()` brings the whole [B, 3] flag table + # over in a single transfer. + flip_rows = None if flips is None else flips.tolist() + out = input.clone() # For each batch element, build list of spatial dims to flip. `input[b]` is # [C, D, H, W], so spatial axis i sits at dim 1 + i. for b in range(batch_size): - if flips is None: + if flip_rows is None: # No sampled flags (a caller invoking apply_transform directly): fall # back to flipping every configured axis. flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis] else: - flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis and bool(flips[b, axis])] + flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis and flip_rows[b][axis]] if len(flip_dims) > 0: out[b] = torch.flip(input[b], dims=tuple(flip_dims)) @@ -610,12 +610,17 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> if len(self.flip_axis) == 0: return {"flip": flips} allowed = torch.as_tensor(self.flip_axis, device=flips.device, dtype=torch.long) - for b in range(batch_size): - if flips[b, allowed].sum() == 0: - # pick one allowed axis at random - choice = int(torch.randint(low=0, high=len(self.flip_axis), size=(1,)).item()) - axis = int(self.flip_axis[choice]) - flips[b, axis] = 1 + # Which samples drew nothing, and the axis to rescue each with, decided for + # the whole batch at once. The loop this replaces read `flips[b, allowed].sum()` + # back to the host for every sample and then drew a separate scalar for the + # ones that needed it. + empty = flips.index_select(1, allowed).sum(dim=1) == 0 + rescue = allowed[torch.randint(low=0, high=len(self.flip_axis), size=(batch_size,), device=flips.device)] + flips.scatter_( + 1, + rescue.view(-1, 1), + torch.where(empty, torch.ones_like(empty, dtype=flips.dtype), flips.gather(1, rescue.view(-1, 1)).squeeze(1)).view(-1, 1), + ) return {"flip": flips} @@ -666,18 +671,24 @@ def __init__( def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: return self.identity_matrix(input) - def _fill_value(self, volume: Tensor, flags: dict[str, Any]) -> float: - """What goes outside the kept box, for one batch element. + def _fill_values(self, batch: Tensor, flags: dict[str, Any]) -> Tensor: + """What goes outside the kept box, as a `[B, 1, 1, 1, 1]` per-sample value. A mask is padded with 0 -- background -- whatever `pad_value` says: a label number is not an intensity, and an argmax over a one-hot volume still has to resolve to the background channel out there. The mask pass identifies itself through `flags["data_keys"]`, which `MaskSequentialOpsCustom` sets before calling `transform_masks` (see gpu/base.py). + + Per-sample and on the device: `float(volume.amin())` once per sample was a + host synchronisation per sample, for a number that never leaves the GPU. """ + shape = (batch.shape[0], *([1] * (batch.dim() - 1))) if DataKey.MASK in (flags or {}).get("data_keys", ()): - return 0.0 - return float(volume.amin()) if self.pad_value == "min" else float(self.pad_value) + return torch.zeros(shape, device=batch.device, dtype=batch.dtype) + if self.pad_value == "min": + return batch.amin(dim=tuple(range(1, batch.dim())), keepdim=True) + return torch.full(shape, float(self.pad_value), device=batch.device, dtype=batch.dtype) @torch.no_grad() def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -694,53 +705,39 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ crops = params["crop"] # shape [B, 3] pos = params["pos"] # shape [B, 3] - # Process per-channel and per-batch element - out = input.clone() + # One transfer for the whole batch's draw, then the identical box arithmetic + # in Python. It used to read six scalars off the device per sample -- each + # `float(...)` on a device tensor blocks the host until the queue drains -- + # and the arithmetic is cheap enough that moving it onto the GPU costs more + # in kernel launches than the transfers it saves. + draws = torch.stack([crops, pos], dim=1).tolist() # [B][2][3], each (x, y, z) + # The discarded region's fill value stays a device tensor -- see the class + # docstring for why it is not 0 -- so the `amin` behind it is one batched + # reduction rather than one host read per sample. + fills = self._fill_values(input, flags) + + out = torch.empty_like(input) for b in range(batch_size): - x = input[b] # [C, D, H, W] + (cx, cy, cz), (px, py, pz) = draws[b] - # determine crop fraction and crop size on the image - cx, cy, cz = crops[b] # interpret crop as fraction of upsampled size to keep - crop_D = max(1, round(float(cz) * D)) - crop_H = max(1, round(float(cy) * H)) - crop_W = max(1, round(float(cx) * W)) - - # determine pos fraction of the image - px, py, pz = pos[b] - - # center position - center_z = float(pz) * D - center_y = float(py) * H - center_x = float(px) * W - - # choose top-left-front corner - start_z = round(center_z - crop_D / 2.0) - start_y = round(center_y - crop_H / 2.0) - start_x = round(center_x - crop_W / 2.0) - - # clamp to valid limits - max_z = max(0, D - crop_D) - max_y = max(0, H - crop_H) - max_x = max(0, W - crop_W) - - z1 = max(0, min(start_z, max_z)) - y1 = max(0, min(start_y, max_y)) - x1 = max(0, min(start_x, max_x)) - - z2 = z1 + crop_D - y2 = y1 + crop_H - x2 = x1 + crop_W - - patch = x[:, z1:z2, y1:y2, x1:x2] - - # Place the patch back where it came from, on a canvas of the discarded - # region's fill value -- see the class docstring for why that value is not 0. - canvas = torch.full((C, D, H, W), self._fill_value(x, flags), dtype=patch.dtype, device=patch.device) - canvas[:, z1:z2, y1:y2, x1:x2] = patch - - out[b] = canvas + crop_D = max(1, round(cz * D)) + crop_H = max(1, round(cy * H)) + crop_W = max(1, round(cx * W)) + + # top-left-front corner of the box centred on (pz, py, px), clamped inside + z1 = max(0, min(round(pz * D - crop_D / 2.0), max(0, D - crop_D))) + y1 = max(0, min(round(py * H - crop_H / 2.0), max(0, H - crop_H))) + x1 = max(0, min(round(px * W - crop_W / 2.0), max(0, W - crop_W))) + + # Fill, then copy the kept box back over it. The old code built a separate + # per-sample canvas and copied that into an `input.clone()`, so the batch + # was written three times over. + out[b] = fills[b] + out[b][:, z1 : z1 + crop_D, y1 : y1 + crop_H, x1 : x1 + crop_W] = input[b][ + :, z1 : z1 + crop_D, y1 : y1 + crop_H, x1 : x1 + crop_W + ] return out diff --git a/smauglab/transforms/gpu/transforms_list.py b/smauglab/transforms/gpu/transforms_list.py index 0d7bfc4..190d32c 100644 --- a/smauglab/transforms/gpu/transforms_list.py +++ b/smauglab/transforms/gpu/transforms_list.py @@ -99,9 +99,17 @@ def _apply_mix(self, x: Tensor, seg: Tensor | None) -> Tensor: if seg is not None: child_params["seg"] = seg - for j in idx.tolist(): + chosen = idx.tolist() + # One read of k uniforms instead of one `if` on a device scalar per + # transform. Comparing a CUDA tensor against `t.p` inside an `if` forces a + # host synchronisation each time, and this runs once per sample per bucket. + # Each transform still gets its own independent draw; they are no longer + # drawn one at a time, so the stream differs from the previous release's. + draws = torch.rand(len(chosen), device=x.device, dtype=x.dtype).tolist() + + for draw, j in zip(draws, chosen): t = self.transforms_list[j] - if torch.rand(1, device=x.device, dtype=x.dtype) > t.p: + if draw > t.p: continue if not hasattr(t, "apply_transform"): raise TypeError(f"All transforms must implement apply_transform like ImageOnlyTransform. Got {type(t)}") @@ -127,14 +135,13 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ return self._apply_mix(input, seg) batch_size = input.shape[0] - # A clone, not `out = input`: the loop writes back through `out[i:i+1]`, so - # without it the caller's batch is modified in place. Every sibling transform - # in gpu/spatial.py clones. - out = input.clone() + # The per-sample results are concatenated rather than written back into a + # clone of the input: the clone wrote the whole batch out only for every row + # of it to be overwritten again. Nothing here writes through `input` -- each + # leaf transform clones before it assigns, which `test_no_inplace_mutation` + # pins -- so the caller's batch is still untouched. + rows = [] for i in range(batch_size): - xi = out[i : i + 1] - seg_i = None seg_i = seg[i : i + 1] if seg is not None and isinstance(seg, torch.Tensor) and seg.shape[0] == batch_size else seg - xi = self._apply_mix(xi, seg_i) - out[i : i + 1] = xi - return out + rows.append(self._apply_mix(input[i : i + 1], seg_i)) + return torch.cat(rows, dim=0) diff --git a/smauglab/transforms/kernels.py b/smauglab/transforms/kernels.py index 242ec07..3eab526 100644 --- a/smauglab/transforms/kernels.py +++ b/smauglab/transforms/kernels.py @@ -23,6 +23,7 @@ from __future__ import annotations +import functools import math from typing import Union @@ -35,12 +36,16 @@ "LAPLACE_3D", "SCHARR_2D", "SCHARR_3D", + "box_filter", + "depthwise_conv3d", "gaussian_blur3d", "gaussian_kernel1d", "gaussian_kernel3d", "laplace_kernel", + "laplacian_response", "random_bias_field3d", "scharr_kernels", + "stacked_scharr_kernels", ] @@ -170,6 +175,7 @@ def gaussian_blur3d( jitter = (1.0 / blur_range) + torch.rand(3, device=device) * (blur_range - 1.0 / blur_range) sigmas = sigmas * jitter + batch = image.shape[0] out = image for axis, s in enumerate(sigmas.tolist()): if s <= 0: @@ -182,10 +188,51 @@ def gaussian_blur3d( shape = [1, 1, 1, 1, 1] shape[2 + axis] = ksize - weight = kernel.view(shape).repeat(channels, 1, 1, 1, 1) + out = depthwise_conv3d(_pad_axis(out, axis, pad, padding_mode), kernel.view(shape)) + return out.view(batch, channels, *out.shape[2:]) + + +def depthwise_conv3d(volume: Tensor, kernel: Tensor, filters_per_plane: int = 1) -> Tensor: + """Convolve every (sample, channel) plane of a pre-padded volume, in one call. + + `volume` is `[B, C, D, H, W]`, already padded. `kernel` is + `[filters_per_plane, 1, kd, kh, kw]` to apply the same filter bank to every + plane, or `[B*C*filters_per_plane, 1, ...]` to give each plane its own. The + result is `[B*C, filters_per_plane, ...]`. + + Folding the batch into the channel axis is the whole point. `F.conv3d` over a + single-channel volume (`groups == 1`) goes through cuDNN's implicit-GEMM path, + which on an A40 takes **4.25 ms** for one 3x3x3 kernel over a 128^3 patch; + with `groups == B*C > 1` it dispatches to ATen's depthwise kernel and the same + convolution takes **0.12 ms**. The depthwise path is also the more accurate of + the two -- implicit GEMM runs in TF32 by default and lands 6e-3 from a float64 + reference where the depthwise kernel lands 4e-6 away. + + A genuinely single-plane call (one sample, one channel, which is what the + per-sample `RandomChooseXTransformsGPU` bucket hands down) cannot reach that + path, so the plane is duplicated to make `groups == 2` and the spare output + dropped. Twice the arithmetic, a twentieth of the time. + """ + if volume.dim() != 5: + raise ValueError(f"expected a 5D [B, C, D, H, W] tensor, got shape {tuple(volume.shape)}") + planes = volume.shape[0] * volume.shape[1] + merged = volume.reshape(1, planes, *volume.shape[2:]) + + if kernel.shape[0] == filters_per_plane: + weight = kernel.repeat(planes, 1, 1, 1, 1) + elif kernel.shape[0] == planes * filters_per_plane: + weight = kernel + else: + raise ValueError(f"kernel must hold {filters_per_plane} or {planes * filters_per_plane} filters, got {kernel.shape[0]}") - out = F.conv3d(_pad_axis(out, axis, pad, padding_mode), weight, groups=channels) - return out + groups = planes + if planes == 1: + merged = merged.expand(1, 2, *merged.shape[2:]).contiguous() + weight = weight.repeat(2, 1, 1, 1, 1) + groups = 2 + + out = F.conv3d(merged, weight, groups=groups) + return out[0, : planes * filters_per_plane].reshape(planes, filters_per_plane, *out.shape[2:]) # --- smooth random fields --------------------------------------------------------- @@ -275,19 +322,85 @@ def random_bias_field3d( ] +# These tables are compile-time constants, but `torch.tensor(nested_list)` walks the +# list in Python and then copies host memory to the device -- every call, from inside +# `apply_transform`, for a tensor of at most 27 numbers. Caching makes the second and +# later calls a dictionary lookup. +# +# The cached tensors are shared, so callers must treat them as read-only. Every one in +# this repository does: they expand, reshape, or convolve with them. + + +@functools.cache +def _laplace_table(spatial_dims: int, device: torch.device | None, dtype: torch.dtype) -> Tensor: + table = LAPLACE_2D if spatial_dims == 2 else LAPLACE_3D + return torch.tensor(table, dtype=dtype, device=device) + + +@functools.cache +def _scharr_table(spatial_dims: int, index: int, device: torch.device | None, dtype: torch.dtype) -> Tensor: + table = SCHARR_2D if spatial_dims == 2 else SCHARR_3D + return torch.tensor(table[index], dtype=dtype, device=device) + + +def box_filter(volume: Tensor, spatial_dims: int) -> Tensor: + """Sum over a 3-wide window on every spatial axis, zero-padded outside. + + Separable, so `spatial_dims` one-dimensional passes -- `3 * d` reads per voxel + instead of `3 ** d`. + """ + conv = F.conv3d if spatial_dims == 3 else F.conv2d + out = volume + for axis in range(spatial_dims): + shape = [1, 1] + [1] * spatial_dims + shape[2 + axis] = 3 + padding = [0] * (2 * spatial_dims) + padding[(spatial_dims - 1 - axis) * 2] = 1 + padding[(spatial_dims - 1 - axis) * 2 + 1] = 1 + out = conv(F.pad(out, padding), torch.ones(shape, device=volume.device, dtype=volume.dtype)) + return out + + +def laplacian_response(volume: Tensor, spatial_dims: int) -> Tensor: + """Convolve `[1, 1, *spatial]` with the Laplacian, without the dense kernel. + + `LAPLACE_3D` is `27 * delta - ones(3, 3, 3)` (and `LAPLACE_2D` is + `9 * delta - ones(3, 3)`), convolution is linear, and the all-ones window is + separable -- so the whole filter is a box filter subtracted from a scaled copy. + That is nine multiply-free passes in 3-D rather than twenty-seven taps per + voxel, and measured 8.6 ms against 32.9 ms for a 96^3 patch on eight CPU + threads. Zero padding, matching `padding="same"`. + + Float addition is not associative, so the result differs from the dense + convolution in the last couple of digits (2e-5 absolute on a response of + magnitude 100). + """ + return (3**spatial_dims) * volume - box_filter(volume, spatial_dims) + + def laplace_kernel(spatial_dims: int, device: torch.device | None = None, dtype: torch.dtype = torch.float32) -> Tensor: - """The Laplacian for 2-D or 3-D data.""" - if spatial_dims == 2: - return torch.tensor(LAPLACE_2D, dtype=dtype, device=device) - if spatial_dims == 3: - return torch.tensor(LAPLACE_3D, dtype=dtype, device=device) - raise ValueError(f"Laplace kernel is defined for 2 or 3 spatial dimensions, got {spatial_dims}") + """The Laplacian for 2-D or 3-D data. Cached and shared; do not mutate the result.""" + if spatial_dims not in (2, 3): + raise ValueError(f"Laplace kernel is defined for 2 or 3 spatial dimensions, got {spatial_dims}") + return _laplace_table(spatial_dims, device, dtype) def scharr_kernels(spatial_dims: int, device: torch.device | None = None, dtype: torch.dtype = torch.float32) -> list[Tensor]: - """The directional Scharr kernels: two for 2-D data, three for 3-D.""" - if spatial_dims == 2: - return [torch.tensor(k, dtype=dtype, device=device) for k in SCHARR_2D] - if spatial_dims == 3: - return [torch.tensor(k, dtype=dtype, device=device) for k in SCHARR_3D] - raise ValueError(f"Scharr kernels are defined for 2 or 3 spatial dimensions, got {spatial_dims}") + """The directional Scharr kernels: two for 2-D data, three for 3-D. + + Cached and shared; do not mutate the results. + """ + if spatial_dims not in (2, 3): + raise ValueError(f"Scharr kernels are defined for 2 or 3 spatial dimensions, got {spatial_dims}") + count = 2 if spatial_dims == 2 else 3 + return [_scharr_table(spatial_dims, i, device, dtype) for i in range(count)] + + +@functools.cache +def stacked_scharr_kernels(spatial_dims: int, device: torch.device | None = None, dtype: torch.dtype = torch.float32) -> Tensor: + """The Scharr kernels as one `[n_dims, 1, *kernel_shape]` conv weight. + + Lets a caller run all the directional filters in a single convolution instead of + one per direction. Cached and shared; do not mutate the result. + """ + return torch.stack(scharr_kernels(spatial_dims, device, dtype)).unsqueeze(1) diff --git a/smauglab/transforms/synthseg/functional.py b/smauglab/transforms/synthseg/functional.py index ab4dd8e..b01d041 100644 --- a/smauglab/transforms/synthseg/functional.py +++ b/smauglab/transforms/synthseg/functional.py @@ -26,6 +26,7 @@ from __future__ import annotations +import functools import math from collections.abc import Sequence from typing import Union @@ -41,6 +42,7 @@ "bias_field", "blurring_sigma_for_downsampling", "convert_labels", + "distinct_values", "em_subdivide_labels", "flip_lr_with_swap", "gaussian_blur_3d", @@ -219,15 +221,28 @@ def sample_gmm_parameters( # Background special-casing (build_model_inputs): per subject, 5% pure black, # 25% very dark/low-variance, 70% normal draw. if randomise_background and background_label_index is not None and 0 <= background_label_index < n_labels: - for b in range(batch): - r = float(torch.rand((), device=device)) - if r > 0.95: - means_lab[b, background_label_index, :] = 0.0 - stds_lab[b, background_label_index, :] = 0.0 - elif r > 0.70: - means_lab[b, background_label_index, :] = torch.rand(n_channels, device=device) * 15.0 - stds_lab[b, background_label_index, :] = torch.rand(n_channels, device=device) * 5.0 - # else: keep the normal draw + # One draw per sample and two candidate parameter sets, selected with + # `torch.where`. The loop this replaces read `float(torch.rand(()))` back to + # the host once per sample, which stalls the whole queue; the branch + # probabilities and the per-sample independence are unchanged, but the dark + # draws are now made for every sample rather than only the ones that take + # that branch, so the stream differs from the previous release's. + draw = torch.rand(batch, 1, device=device) + dark_means = torch.rand(batch, n_channels, device=device) * 15.0 + dark_stds = torch.rand(batch, n_channels, device=device) * 5.0 + black = draw > 0.95 + dark = (draw > 0.70) & ~black + + background_means = torch.where( + black, torch.zeros_like(dark_means), torch.where(dark, dark_means, means_lab[:, background_label_index, :]) + ) + background_stds = torch.where( + black, torch.zeros_like(dark_stds), torch.where(dark, dark_stds, stds_lab[:, background_label_index, :]) + ) + means_lab = means_lab.clone() + stds_lab = stds_lab.clone() + means_lab[:, background_label_index, :] = background_means + stds_lab[:, background_label_index, :] = background_stds return means_lab, stds_lab @@ -262,14 +277,15 @@ def labels_to_image_gmm( lut[label_values] = torch.arange(label_values.numel(), device=device) idx = lut[label_map.clamp(min=0, max=max_label)].squeeze(1) # (B, D, H, W) - image = torch.empty((B, C, *spatial), device=device, dtype=dtype) - for b in range(B): - idx_b = idx[b] - for ch in range(C): - mean_map = means[b, :, ch][idx_b] - std_map = stds[b, :, ch][idx_b] - image[b, ch] = mean_map + std_map * torch.randn(spatial, device=device, dtype=dtype) - return image + # One gather and one noise draw for the whole batch. The nested loop this + # replaces indexed `means[b, :, ch][idx_b]` once per (sample, channel) and drew + # its noise volume separately, so a four-channel batch of two paid eight small + # `randn` launches and eight gathers for what is one of each. + flat_idx = idx.reshape(B, 1, -1).expand(B, C, -1) # (B, C, N) + mean_map = torch.gather(means.transpose(1, 2), 2, flat_idx) # (B, C, N) + std_map = torch.gather(stds.transpose(1, 2), 2, flat_idx) + noise = torch.randn((B, C, *spatial), device=device, dtype=dtype) + return (mean_map + std_map * noise.reshape(B, C, -1)).reshape(B, C, *spatial) # --------------------------------------------------------------------------- @@ -355,8 +371,16 @@ def stack3(rows): return affine +@functools.lru_cache(maxsize=8) def _identity_grid(shape: tuple[int, int, int], device: torch.device) -> torch.Tensor: - """Voxel-coordinate identity grid in ``(i, j, k)`` order, shape ``(3, D, H, W)``.""" + """Voxel-coordinate identity grid in ``(i, j, k)`` order, shape ``(3, D, H, W)``. + + Cached: it depends only on the patch shape, and `_integrate_velocity` calls + `warp_volume` seven times per field, each of which rebuilt it -- 24 MB of + `meshgrid` and `stack` at 128^3, seven times over, for a constant. The returned + tensor is shared; callers must not write into it (none do -- `warp_volume` + expands and then reads). + """ d, h, w = shape zs = torch.arange(d, device=device, dtype=torch.float32) ys = torch.arange(h, device=device, dtype=torch.float32) @@ -643,41 +667,90 @@ def mimic_acquisition( # --------------------------------------------------------------------------- # EM label completion for sparse label maps (SynthSeg paper, Sec. 5.4) # --------------------------------------------------------------------------- +def _quadratic_log_likelihood_terms(means: torch.Tensor, var: torch.Tensor, weights: torch.Tensor, eps: float) -> torch.Tensor: + """The `(3, K)` coefficients of the per-component log-density as a quadratic in x. + + ``log N(x | m, v) + log w`` expands to ``a x^2 + b x + c`` with + ``a = -1/(2v)``, ``b = m/v`` and ``c = log w - (log 2pi + log v)/2 - m^2/(2v)``. + Stacking those lets a whole `(n, K)` responsibility matrix come out of one + matrix multiply against ``[x^2, x, 1]``, instead of half a dozen elementwise + kernels each writing an `(n, K)` temporary. + """ + inv_var = 1.0 / var + const = torch.log(weights.clamp_min(eps)) - 0.5 * (math.log(2.0 * math.pi) + torch.log(var)) - 0.5 * means * means * inv_var + return torch.stack([-0.5 * inv_var, means * inv_var, const]) + + +def distinct_values(volume: torch.Tensor) -> torch.Tensor: + """The distinct non-negative integer values in `volume`, ascending. + + `torch.unique` radix-sorts the input; counting into bins and keeping the + non-empty ones answers the same question for small label alphabets in about + half the time on a 2 x 128^3 volume, and the two results are identical. + """ + return torch.bincount(volume.reshape(-1).clamp_min(0)).nonzero().flatten() + + +def _label_counts_per_sample(label_map: torch.Tensor) -> torch.Tensor: + """`(B, max_label + 1)` voxel counts per label, per sample.""" + batch = label_map.shape[0] + flat = label_map.reshape(batch, -1).clamp_min(0) + n_bins = int(flat.max()) + 1 + # One `bincount` over row-offset labels: a `scatter_add_` into a handful of bins + # has every thread fighting for the same addresses, where bincount accumulates + # in shared memory per block. + offset = flat + torch.arange(batch, device=flat.device).view(-1, 1) * n_bins + return torch.bincount(offset.reshape(-1), minlength=batch * n_bins).view(batch, n_bins) + + def _em_gmm_1d(x_fit: torch.Tensor, n_components: int, n_iters: int, eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: """Fit a 1D Gaussian mixture by Expectation-Maximization. ``x_fit`` is a 1D tensor of intensities. Returns ``(means, vars, weights)`` each of shape ``(K,)`` (``K = min(n_components, len(x_fit))``), or ``None`` if ``x_fit`` is empty. Responsibilities use a numerically-stable log-sum-exp. + + Each iteration is two matrix multiplies and a softmax rather than a dozen + elementwise passes over an ``(n, K)`` array: the E-step log-densities are a + quadratic in x (see :func:`_quadratic_log_likelihood_terms`) and the M-step only + needs the responsibility-weighted sums of ``1``, ``x`` and ``x^2``, which is the + same design matrix transposed. With ``n_iters=20`` and up to ten components that + is the difference between ~240 kernel launches per fit and ~60. + + ``x`` is centred on its own mean before the moments are accumulated, so the + ``E[x^2] - E[x]^2`` variance cannot lose its significant digits to cancellation + when a component is tight and far from zero. The means are shifted back at the + end. """ n = x_fit.numel() if n == 0: return None k = max(1, min(int(n_components), n)) device = x_fit.device - xmin, xmax = float(x_fit.min()), float(x_fit.max()) - means = torch.linspace(xmin, xmax, k, device=device) - if xmax <= xmin: # degenerate (constant region): spread the means slightly - means = means + torch.arange(k, device=device, dtype=means.dtype) * eps + + # `float(x.min())` twice was two host synchronisations per fit, and + # `em_subdivide_labels` calls this once per label per sample. + x_min, x_max = x_fit.min(), x_fit.max() + step = (x_max - x_min) / max(k - 1, 1) + means = x_min + step * torch.arange(k, device=device, dtype=x_fit.dtype) + # Degenerate (constant region): spread the means slightly. + means = torch.where(x_max > x_min, means, x_min + torch.arange(k, device=device, dtype=x_fit.dtype) * eps) var = x_fit.var(unbiased=False).clamp_min(eps).repeat(k) weights = torch.full((k,), 1.0 / k, device=device) - x = x_fit.view(n, 1) - log2pi = math.log(2.0 * math.pi) + shift = x_fit.mean() + centred = (x_fit - shift).view(n, 1) + design = torch.cat([centred * centred, centred, torch.ones_like(centred)], dim=1) # (n, 3) + means = means - shift + for _ in range(int(n_iters)): - logp = ( - torch.log(weights.clamp_min(eps)).view(1, k) - - 0.5 * (log2pi + torch.log(var).view(1, k)) - - 0.5 * (x - means.view(1, k)) ** 2 / var.view(1, k) - ) - logp = logp - torch.logsumexp(logp, dim=1, keepdim=True) - resp = logp.exp() # (n, k) - nk = resp.sum(0).clamp_min(eps) # (k,) + resp = torch.softmax(design @ _quadratic_log_likelihood_terms(means, var, weights, eps), dim=1) + moments = resp.transpose(0, 1) @ design # (k, 3): sums of r*x^2, r*x, r + nk = moments[:, 2].clamp_min(eps) weights = nk / n - means = (resp * x).sum(0) / nk - var = (resp * (x - means.view(1, k)) ** 2).sum(0) / nk - var = var.clamp_min(eps) - return means, var, weights + means = moments[:, 1] / nk + var = (moments[:, 0] / nk - means * means).clamp_min(eps) + return means + shift, var, weights def _assign_gmm( @@ -690,16 +763,15 @@ def _assign_gmm( ) -> torch.Tensor: """Hard-assign each value in ``x_full`` to its most likely mixture component.""" n = x_full.numel() - k = means.numel() out = torch.empty(n, dtype=torch.long, device=x_full.device) - logw = torch.log(weights.clamp_min(eps)).view(1, k) - half_logvar = 0.5 * torch.log(var).view(1, k) - m = means.view(1, k) - v = var.view(1, k) + # Same quadratic form as the E-step: one matmul per chunk instead of the + # subtract / square / divide / subtract chain, each of which wrote its own + # (chunk, K) array out to memory. + terms = _quadratic_log_likelihood_terms(means, var, weights, eps) for s in range(0, n, chunk): xc = x_full[s : s + chunk].view(-1, 1) - logp = logw - half_logvar - 0.5 * (xc - m) ** 2 / v - out[s : s + chunk] = logp.argmax(dim=1) + design = torch.cat([xc * xc, xc, torch.ones_like(xc)], dim=1) + out[s : s + chunk] = (design @ terms).argmax(dim=1) return out @@ -749,7 +821,14 @@ def em_subdivide_labels( B = image.shape[0] device = image.device ref = image[:, channel] # (B, D, H, W) - parents = torch.unique(label_map).long().tolist() # sorted, batch-wide + # Per-sample label histogram, in one pass. It answers three questions the loop + # below used to ask one at a time with a host synchronisation each: which labels + # exist batch-wide, which is the largest, and how many voxels each holds in each + # sample. + label_counts = _label_counts_per_sample(label_map) + totals = label_counts.sum(dim=0) + parents = totals.nonzero().flatten().tolist() # sorted, batch-wide + counts_per_sample = label_counts.tolist() parent_to_idx = {p: i for i, p in enumerate(parents)} lo, hi = int(background_clusters_range[0]), int(background_clusters_range[1]) mult = max(hi, int(n_foreground_clusters)) + 1 # collision-free encoding @@ -759,7 +838,7 @@ def em_subdivide_labels( # fall back to the largest-area label so it still receives the richer # [min, max] background split rather than the 2-cluster foreground split. if background_label not in parents and len(parents) > 0: - background_label = int(torch.bincount(label_map.flatten().clamp_min(0)).argmax()) + background_label = int(totals.argmax()) bg_n_shared = int(torch.randint(lo, hi + 1, (1,), device=device)) if same_on_batch else None @@ -769,10 +848,10 @@ def em_subdivide_labels( ref_b = ref[b] for p in parents: pi = parent_to_idx[p] - mask = lab_b == p - cnt = int(mask.sum()) + cnt = int(counts_per_sample[b][p]) if cnt == 0: continue + mask = lab_b == p if p == background_label: k = bg_n_shared if bg_n_shared is not None else int(torch.randint(lo, hi + 1, (1,), device=device)) else: @@ -792,7 +871,7 @@ def em_subdivide_labels( assign = torch.zeros(cnt, dtype=torch.long, device=device) if fit is None else _assign_gmm(x, *fit, eps=eps) fine[b, 0][mask] = pi * mult + assign - gen_values = torch.unique(fine).long().tolist() + gen_values = distinct_values(fine).tolist() out_values = [parents[g // mult] for g in gen_values] return fine, gen_values, out_values diff --git a/unit_tests/test_fast_paths.py b/unit_tests/test_fast_paths.py new file mode 100644 index 0000000..a173274 --- /dev/null +++ b/unit_tests/test_fast_paths.py @@ -0,0 +1,307 @@ +"""Each optimised primitive, against the straightforward code it replaced. + +The speed work in this area replaced several inner loops with batched or +reformulated equivalents. Every one of them is supposed to compute the same thing +as the obvious spelling, so every one of them is checked against the obvious +spelling here rather than against a stored expected value -- that way the test +says what the optimisation claims, and keeps saying it if the fast path is ever +rewritten again. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +from smauglab.transforms.gpu.contrast import ( + RandomBiasFieldGPU, + _axis_power_tables, + _batched_quantiles, + apply_convolution_per_sample, +) +from smauglab.transforms.gpu.fromSeg import _kmeans_1d, _region_moments, foreground_classes, segment_sum, voxel_coordinates +from smauglab.transforms.kernels import LAPLACE_2D, LAPLACE_3D, depthwise_conv3d, laplace_kernel, laplacian_response +from smauglab.transforms.synthseg.functional import _assign_gmm, _em_gmm_1d, distinct_values +from unit_tests.helpers import SmaugLabTestCase + +CPU = torch.device("cpu") + + +class TestDepthwiseConv3d(SmaugLabTestCase): + """`depthwise_conv3d` folds the batch into `groups`; the result must not move.""" + + @staticmethod + def _reference(padded: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor: + """One `F.conv3d` per (sample, channel) plane -- what the loops used to do.""" + planes = padded.shape[0] * padded.shape[1] + flat = padded.reshape(planes, 1, *padded.shape[2:]) + return torch.stack( + [F.conv3d(flat[p : p + 1], kernel[min(p, kernel.shape[0] - 1) : min(p, kernel.shape[0] - 1) + 1])[0] for p in range(planes)] + ) + + def test_a_shared_kernel_matches_a_per_plane_loop(self): + for batch, channels in ((1, 1), (3, 1), (2, 2)): + with self.subTest(batch=batch, channels=channels): + volume = torch.randn(batch, channels, 7, 8, 9) + kernel = torch.randn(1, 1, 3, 3, 3) + padded = F.pad(volume, [1] * 6, mode="reflect") + got = depthwise_conv3d(padded, kernel) + self.assertTrue(torch.allclose(got, self._reference(padded, kernel), atol=1e-5)) + + def test_the_single_plane_duplication_does_not_change_the_answer(self): + """A one-sample, one-channel call is the case that needs `groups == 2` faked.""" + volume = torch.randn(1, 1, 6, 7, 8) + kernel = torch.randn(1, 1, 3, 3, 3) + padded = F.pad(volume, [1] * 6, mode="reflect") + + got = depthwise_conv3d(padded, kernel) + + self.assertEqual(got.shape, (1, 1, 6, 7, 8)) + self.assertTrue(torch.allclose(got, F.conv3d(padded, kernel), atol=1e-5)) + + def test_a_per_plane_kernel_bank_is_applied_plane_by_plane(self): + volume = torch.randn(3, 1, 6, 6, 6) + kernels = torch.randn(3, 1, 3, 3, 3) + padded = F.pad(volume, [1] * 6, mode="reflect") + + got = depthwise_conv3d(padded, kernels) + + for plane in range(3): + expected = F.conv3d(padded[plane : plane + 1], kernels[plane : plane + 1]) + self.assertTrue(torch.allclose(got[plane : plane + 1], expected, atol=1e-5)) + + def test_several_filters_per_plane_land_on_the_channel_axis(self): + volume = torch.randn(2, 1, 6, 6, 6) + bank = torch.randn(4, 1, 3, 3, 3) + padded = F.pad(volume, [1] * 6, mode="reflect") + + got = depthwise_conv3d(padded, bank, filters_per_plane=4) + + self.assertEqual(got.shape, (2, 4, 6, 6, 6)) + for plane in range(2): + expected = F.conv3d(padded[plane : plane + 1], bank) + self.assertTrue(torch.allclose(got[plane : plane + 1], expected, atol=1e-5)) + + +class TestPerSampleConvolution(SmaugLabTestCase): + def test_it_matches_convolving_each_sample_separately(self): + volume = torch.randn(3, 6, 7, 8) + kernels = [torch.randn(3, 3, 3) for _ in range(3)] + + got = apply_convolution_per_sample(volume, kernels) + + for b, kernel in enumerate(kernels): + padded = F.pad(volume[b : b + 1].unsqueeze(1), [1] * 6, mode="reflect") + expected = F.conv3d(padded, kernel.view(1, 1, 3, 3, 3))[0, 0] + self.assertTrue(torch.allclose(got[b], expected, atol=1e-5)) + + def test_kernels_of_different_sizes_are_zero_padded_not_resized(self): + """A 1^3 kernel padded out to 7^3 must still be a 1^3 convolution.""" + volume = torch.randn(2, 10, 10, 10) + small = torch.full((1, 1, 1), 2.0) + large = torch.randn(7, 7, 7) + + got = apply_convolution_per_sample(volume, [small, large]) + + self.assertTrue(torch.allclose(got[0], volume[0] * 2.0, atol=1e-5)) + padded = F.pad(volume[1:2].unsqueeze(1), [3] * 6, mode="reflect") + self.assertTrue(torch.allclose(got[1], F.conv3d(padded, large.view(1, 1, 7, 7, 7))[0, 0], atol=1e-4)) + + +class TestLaplacianIdentity(SmaugLabTestCase): + """`laplacian_response` uses `3**d * x - box(x)` instead of the dense kernel.""" + + def test_it_matches_the_dense_kernel_in_three_dimensions(self): + volume = torch.randn(1, 1, 9, 10, 11) + dense = F.conv3d(volume, laplace_kernel(3).view(1, 1, 3, 3, 3), padding="same") + self.assertTrue(torch.allclose(laplacian_response(volume, 3), dense, atol=1e-4)) + + def test_it_matches_the_dense_kernel_in_two_dimensions(self): + image = torch.randn(1, 1, 11, 12) + dense = F.conv2d(image, laplace_kernel(2).view(1, 1, 3, 3), padding="same") + self.assertTrue(torch.allclose(laplacian_response(image, 2), dense, atol=1e-4)) + + def test_the_tables_really_are_a_scaled_delta_minus_a_box(self): + """The identity the reformulation rests on, stated directly.""" + for dims, table in ((2, LAPLACE_2D), (3, LAPLACE_3D)): + with self.subTest(dims=dims): + kernel = torch.tensor(table, dtype=torch.float32) + box = torch.ones_like(kernel) + delta = torch.zeros_like(kernel) + delta[(1,) * dims] = float(3**dims) + self.assertTrue(torch.equal(kernel, delta - box)) + + +class TestCachedKernels(SmaugLabTestCase): + def test_repeated_calls_return_the_same_tensor(self): + self.assertIs(laplace_kernel(3, CPU), laplace_kernel(3, CPU)) + + def test_the_cached_kernel_still_holds_the_table(self): + self.assertTrue(torch.equal(laplace_kernel(3, CPU), torch.tensor(LAPLACE_3D, dtype=torch.float32))) + + +class TestSegmentSum(SmaugLabTestCase): + def test_it_matches_scatter_add(self): + ids = torch.randint(0, 6, (5000,)) + values = torch.rand(5000) + expected = torch.zeros(6).scatter_add_(0, ids, values) + self.assertTrue(torch.allclose(segment_sum(ids, values, 6), expected, atol=1e-4)) + + def test_a_trailing_empty_segment_is_still_returned(self): + """`bincount` sizes itself from the largest id present; the pad restores it.""" + ids = torch.zeros(10, dtype=torch.long) + got = segment_sum(ids, torch.ones(10), 4) + self.assertEqual(got.shape, (4,)) + self.assertTrue(torch.equal(got, torch.tensor([10.0, 0.0, 0.0, 0.0]))) + + +class TestDistinctValues(SmaugLabTestCase): + def test_foreground_classes_matches_unique(self): + labels = torch.randint(0, 5, (2, 1, 6, 6, 6)) + expected = labels.unique() + self.assertTrue(torch.equal(foreground_classes(labels), expected[expected > 0])) + + def test_distinct_values_matches_unique(self): + volume = torch.randint(0, 9, (2, 1, 5, 5, 5)) + self.assertTrue(torch.equal(distinct_values(volume), volume.unique())) + + +class TestBatchedQuantiles(SmaugLabTestCase): + def test_it_matches_torch_quantile_per_sample(self): + data = torch.randn(4, 1, 7, 8, 9) + lower = torch.tensor([0.0, 0.05, 0.2, 0.33]) + upper = torch.tensor([1.0, 0.95, 0.8, 0.67]) + + got_lo, got_hi = _batched_quantiles(data, lower, upper) + + for b in range(4): + flat = data[b].flatten() + self.assertAlmostEqual(float(got_lo[b]), float(torch.quantile(flat, lower[b])), places=5) + self.assertAlmostEqual(float(got_hi[b]), float(torch.quantile(flat, upper[b])), places=5) + + def test_the_no_op_shortcut_returns_the_actual_range(self): + data = torch.randn(3, 1, 5, 5, 5) + lo, hi = _batched_quantiles(data, torch.zeros(3), torch.ones(3)) + self.assertTrue(torch.equal(lo, data.reshape(3, -1).amin(dim=1))) + self.assertTrue(torch.equal(hi, data.reshape(3, -1).amax(dim=1))) + + +class TestRegionMoments(SmaugLabTestCase): + def test_it_matches_the_explicit_masked_statistics(self): + values = torch.rand(400) + masks = torch.rand(5, 400) > 0.6 + design = torch.stack([torch.ones_like(values), values, values * values], dim=1) + + counts, means, stds = _region_moments(masks, design) + + for r in range(5): + selected = values[masks[r]] + self.assertAlmostEqual(float(counts[r]), max(float(masks[r].sum()), 1.0), places=4) + self.assertAlmostEqual(float(means[r]), float(selected.mean()), places=4) + self.assertAlmostEqual(float(stds[r]), float(selected.std(unbiased=False)), places=4) + + +class TestKMeans1D(SmaugLabTestCase): + def test_it_separates_two_well_separated_clusters(self): + values = torch.cat([torch.randn(500) * 0.02, torch.randn(500) * 0.02 + 1.0]) + centroids = _kmeans_1d(values, 2) + low, high = sorted(float(c) for c in centroids) + self.assertAlmostEqual(low, 0.0, places=1) + self.assertAlmostEqual(high, 1.0, places=1) + + def test_the_centroids_stay_sorted(self): + """The `searchsorted` assignment is only valid while they are.""" + centroids = _kmeans_1d(torch.rand(2000), 5) + self.assertTrue(bool((centroids[1:] >= centroids[:-1]).all())) + + def test_a_constant_input_does_not_produce_nan(self): + centroids = _kmeans_1d(torch.full((100,), 0.5), 4) + self.assertTrue(bool(torch.isfinite(centroids).all())) + + +class TestBiasFieldContraction(SmaugLabTestCase): + """The polynomial is evaluated as a separable contraction, not a term loop.""" + + def test_it_matches_the_explicit_monomial_sum(self): + for order in (0, 1, 3): + with self.subTest(order=order): + transform = RandomBiasFieldGPU(order=order, p=1.0) + spatial = (5, 6, 7) + coefficients = torch.randn(transform._num_coeffs(3), 2) + + cube = transform._coefficient_cube(coefficients, 3) + x_pow, y_pow, z_pow = _axis_power_tables(order, spatial, CPU, torch.float32) + partial = torch.einsum("bxyz,zd->bxyd", cube, z_pow) + partial = torch.einsum("bxyd,yh->bxdh", partial, y_pow) + got = torch.einsum("bxdh,xw->bdhw", partial, x_pow) + + grids = transform._make_grids(spatial, CPU, torch.float32) + expected = torch.zeros((2, *spatial)) + index = 0 + for xo in range(order + 1): + for yo in range(order + 1 - xo): + for zo in range(order + 1 - (xo + yo)): + term = grids[0].pow(xo) * grids[1].pow(yo) * grids[2].pow(zo) + expected += coefficients[index].view(-1, 1, 1, 1) * term + index += 1 + self.assertTrue(torch.allclose(got, expected, atol=1e-4)) + + +class TestEmGmm(SmaugLabTestCase): + """The EM steps are matrix multiplies now; the fit itself must still be a fit.""" + + def test_it_recovers_two_separated_components(self): + data = torch.cat([torch.randn(800) * 0.1 - 2.0, torch.randn(800) * 0.1 + 3.0]) + means, var, weights = _em_gmm_1d(data, 2, 40, 1e-6) + low, high = sorted(float(m) for m in means) + self.assertAlmostEqual(low, -2.0, places=1) + self.assertAlmostEqual(high, 3.0, places=1) + self.assertTrue(bool((var > 0).all())) + self.assertAlmostEqual(float(weights.sum()), 1.0, places=4) + + def test_the_hard_assignment_splits_them_the_obvious_way(self): + data = torch.cat([torch.randn(400) * 0.1 - 2.0, torch.randn(400) * 0.1 + 3.0]) + fit = _em_gmm_1d(data, 2, 40, 1e-6) + assignment = _assign_gmm(data, *fit, eps=1e-6) + self.assertEqual(len(assignment.unique()), 2) + self.assertTrue(bool((assignment[:400] == assignment[0]).all())) + self.assertTrue(bool((assignment[400:] == assignment[-1]).all())) + self.assertNotEqual(int(assignment[0]), int(assignment[-1])) + + def test_a_constant_region_stays_finite(self): + means, var, weights = _em_gmm_1d(torch.full((200,), 1.5), 3, 20, 1e-6) + for name, tensor in (("means", means), ("var", var), ("weights", weights)): + self.assertTrue(bool(torch.isfinite(tensor).all()), f"{name} went non-finite") + + def test_the_quadratic_form_is_the_gaussian_log_density(self): + """`a x^2 + b x + c` has to be `log w + log N(x | m, v)`, or the E-step is wrong.""" + from smauglab.transforms.synthseg.functional import _quadratic_log_likelihood_terms + + means = torch.tensor([0.5, -1.25]) + var = torch.tensor([0.3, 2.0]) + weights = torch.tensor([0.4, 0.6]) + x = torch.linspace(-3, 3, 11).view(-1, 1) + + terms = _quadratic_log_likelihood_terms(means, var, weights, 1e-6) + got = torch.cat([x * x, x, torch.ones_like(x)], dim=1) @ terms + + expected = ( + torch.log(weights).view(1, 2) + - 0.5 * (math.log(2.0 * math.pi) + torch.log(var).view(1, 2)) + - 0.5 * (x - means.view(1, 2)) ** 2 / var.view(1, 2) + ) + self.assertTrue(torch.allclose(got, expected, atol=1e-5)) + + +class TestVoxelCoordinateCache(SmaugLabTestCase): + def test_the_grid_is_the_meshgrid_it_replaced(self): + coords = voxel_coordinates((2, 3, 4), CPU) + expected = torch.stack(torch.meshgrid(torch.arange(2.0), torch.arange(3.0), torch.arange(4.0), indexing="ij"), dim=-1).reshape( + 24, 3 + ) + self.assertTrue(torch.equal(coords, expected)) + + def test_the_same_shape_is_not_rebuilt(self): + self.assertIs(voxel_coordinates((4, 4, 4), CPU), voxel_coordinates((4, 4, 4), CPU)) diff --git a/unit_tests/test_region_and_stats.py b/unit_tests/test_region_and_stats.py index e08e7f0..896aff1 100644 --- a/unit_tests/test_region_and_stats.py +++ b/unit_tests/test_region_and_stats.py @@ -182,10 +182,14 @@ class TestNonFiniteGuard(SmaugLabTestCase): and the loop assigns into `channel_data[b]` -- so by the time the non-finite guard at the bottom `continue`d, every value it meant to withhold was already in the batch, and the closing `input[:, c] = channel_data` was a no-op. With the clone the - guard does what it says. Reaching it needs the equalisation to *produce* a - non-finite value from finite input, which no input constructed here manages: a NaN - supplied by the caller raises out of `torch.histc` first. So these tests pin the - behaviour the clone must not disturb, and pass either way. + guard does what it says. + + The guard is now also reachable. The equalisation used to call `torch.histc`, + which rejects a non-finite range before any of this runs, so a NaN handed in by + the caller came back as a RuntimeError out of the middle of a training step. The + batched version derives the bin edges itself, so a NaN reaches the guard and the + guard does what it is for: the channel is returned as it came in, with a warning + naming the transform. """ def test_a_degenerate_constant_channel_stays_finite(self): @@ -197,14 +201,18 @@ def test_a_degenerate_constant_channel_stays_finite(self): self.assertTrue(bool(torch.isfinite(out).all()), "a non-finite result reached the batch") - def test_a_non_finite_input_is_still_rejected_loudly(self): - """Documents why the guard cannot be exercised: histc rejects the range first.""" + def test_a_non_finite_input_is_caught_by_the_guard(self): + """A NaN from the caller leaves the channel alone instead of raising.""" transform = RandomHistogramEqualizationGPU(p=1.0, mix_prob=0.0) volume = torch.rand(1, 1, 8, 8, 8) volume[0, 0, 0, 0, 0] = float("nan") + reference = volume.clone() - with self.assertRaises(RuntimeError): - transform.apply_transform(volume, {}, {}, transform=None) + out = transform.apply_transform(volume, {}, {}, transform=None) + + finite = torch.isfinite(reference) + self.assertTrue(torch.equal(out[finite], reference[finite]), "the guard rewrote a channel it rejected") + self.assertTrue(bool(torch.isnan(out[0, 0, 0, 0, 0])), "the caller's own NaN should still be there") def test_equalisation_still_changes_the_image(self): transform = RandomHistogramEqualizationGPU(p=1.0, mix_prob=0.0)