Large scans and several GPUs

This page says which parts of a reconstruction must fit in GPU memory, how to reconstruct scans that do not fit, and how to share a scan among several GPUs. Streaming from host memory applies to projections held as NumPy arrays or memmaps, which is what tj.load returns. A JAX array already on the GPU stays there.

GPU allocatorCopy Anchor Link

Importing tomojax defaults JAX to on-demand allocation rather than an eager 75% GPU reservation, leaving room for cuSolver and CuPy workspaces. The pool retains memory for reuse; Python keeps JAX's default pool limit, while GPU CLI commands set a 90% ceiling. Explicit XLA_PYTHON_CLIENT_PREALLOCATE, XLA_CLIENT_MEM_FRACTION (or JAX's legacy XLA_PYTHON_CLIENT_MEM_FRACTION) and XLA_PYTHON_CLIENT_ALLOCATOR take precedence; the align CLI defaults to the releasing platform allocator. These process-wide settings cannot change an already initialized JAX backend: import TomoJAX or set the environment before the first JAX operation in a notebook.

What has to fit on the GPUCopy Anchor Link

stream_projections = true or false in the method's configuration (a --config key) overrides the 40% rule. Streaming needs the batched projection operators: CGLS streams only with the canonical detector grid, and FISTA-TV and SPDHG-TV not with an explicit detector grid or the exact ray integrator. Streamed CGLS solves the equivalent normal equations, which square the condition number of each step's recurrence. It recomputes the exact gradient periodically, as in-core CGLS recomputes its residual.

The iterative solvers need the whole volume and their work arrays on the GPU; a volume larger than the GPU cannot be reconstructed iteratively. The memory measurements give peak use for 512³ laminography on an 8 GB GPU.

Reconstruct a volume larger than the GPUCopy Anchor Link

Cone beamCopy Anchor Link

tj.reconstruct(scan) and tomojax recon (FDK, the default method on cone data) switch to host z-slabs when the volume would take more than 40% of the free GPU memory. The volume then comes back as a NumPy array, and Reconstruction.info["host_slabs"] is true. Each slab filters only the detector rows it projects onto. The tomojax commands still hold the whole projection stack and volume in host RAM; to write the volume straight to disk, pass a memmap through out=:

import numpy as np

import tomojax as tj

scan = tj.load("scan.nxs")
grid = scan.grid
volume = np.lib.format.open_memmap(
    "volume.npy", mode="w+", dtype=np.float32, shape=(grid.nx, grid.ny, grid.nz)
)
recon = tj.reconstruct(scan, "fbp", out=volume)
volume.flush()

For explicit slab-depth control, use fdk_host with FDKHostConfig(slices_per_batch=...); by default slabs take about a third of the free GPU memory. FDK accepts per-view cone geometry, including ConeSegments; noncircular trajectories retain FDK's approximation.

Parallel beam and laminographyCopy Anchor Link

tj.reconstruct(scan, "fbp") and tomojax recon --method fbp switch to host slabs when the volume would take more than 40% of the free device memory. Reconstruction.info["host_slabs"] records that choice. Passing out= forces host slabs even for a small volume and fills the supplied float32 storage without allocating another full volume:

import numpy as np

import tomojax as tj

scan = tj.load("corrected-lamino.nxs")
grid = scan.grid
volume = np.lib.format.open_memmap(
    "volume.npy", mode="w+", dtype=np.float32, shape=(grid.nx, grid.ny, grid.nz)
)
recon = tj.reconstruct(scan, "fbp", filter="ramp", out=volume)
volume.flush()

Built-in parallel scans use balanced axial slabs of at most 16 slices, reading only the detector rows they need. Automatically allocated axial host output uses Fortran order for contiguous xy-planes; laminography keeps C order for x-slabs. For large disk-backed axial FBP output, pass fortran_order=True to open_memmap: contiguous xy-planes avoid repeatedly dirtying pages across the whole file as successive z-slabs are written. This changes storage order, not the volume's (x, y, z) shape or coordinates. Keep the default C order for laminography's x-slabs.

Laminography and posed scans use x-slabs, each streaming and refiltering every view; those slabs use about half the free GPU memory by default. For explicit slab-depth control, call fbp_host with FBPHostConfig(slices_per_batch=...). Host slabs are not differentiable. Device projections remain resident and are sliced on the device; use NumPy arrays or memmaps to stream projections from host memory. Axial CUDA FBP may stage one slab's detector rows (up to 512 MiB and a quarter of free device memory) to reduce host dispatch; filtering still uses bounded view batches. Input and output must not overlap, and public out= does not support explicit detector grids. method="fourier" also accepts out= for axial parallel scans. These storage choices do not make raw-frame correction out of core: frames.corrected() still returns a complete corrected projection stack in host memory. Axial F-order output can overlap one slab download/write with the next computation when spare device memory permits. All writes finish before return; a failure can leave a partly written volume. A 1024³ laminography volume takes about 37 s on an 8 GB laptop GPU.

Stream raw framesCopy Anchor Link

For cone-beam FDK, tj.reconstruct_slabs(frames, "fbp") corrects only the detector-row band needed by each output slab, using bounded device batches. Neither the full corrected stack nor the whole output volume needs to fit in host memory. One corrected row band and any retained blocks must fit in host memory. Angular weights use the complete acquisition, not the row crop.

For an unposed axial parallel scan, tj.reconstruct_slabs corrects raw rows and reconstructs them without storing the full corrected stack. Choose "fbp", "fourier", "sirt" or "mlem", with that method's configuration through config= when needed; warm_start=True starts SIRT and MLEM from FBP. The whole Diamond I12 scan, a 56.6 GB SIRT volume, took 24 minutes with 1.4 GB of GPU memory. For example, write owned blocks directly to plane-contiguous storage:

import numpy as np

import tomojax as tj
from tomojax.corrections import Stripes

frames = tj.load_frames("raw.nxs")
grid = frames.geometry.grid
volume = np.lib.format.open_memmap(
    "volume.npy", mode="w+", dtype=np.float32,
    shape=(grid.nx, grid.ny, grid.nz), fortran_order=True,
)
with tj.reconstruct_slabs(frames, "fbp", corrections=[Stripes(), Stripes(51, varying=True)]) as slabs:
    for z_slice, block in slabs:
        volume[:, :, z_slice] = block.volume
volume.flush()

Each block retains its global physical grid, genuinely corrected source rows (including interpolation neighbors), and correction records. Blocks remain valid after iteration or context exit, but collecting them retains both their volumes and projections. The context drains pending work on early exit or failure. Only whole-row line-integral corrections are supported, not Paganin or view rejection; keep the raw and calibration inputs unchanged in the context.

Bin detectors finer than the gridCopy Anchor Link

When the detector samples the rotation axis more finely than the voxels (the pixel pitch, divided by the magnification in a cone beam, is smaller than a voxel), the projectors do work the volume cannot use. scan.binned(factor) averages factor × factor pixels, which cuts projection cost by factor². Pixels that do not fill a block at the edges are dropped, and the detector centre moves to match. Iterative tj.reconstruct and tj.align warn, with the factor to use, when the detector is at least twice as fine as the grid. Binning is a Python step; the commands have no binning option. On the FIPS walnut, binning 2 × 2 made least squares 3.5 times faster for no loss of accuracy.

For alignment, stopping the coarse-to-fine schedule early has a similar effect. See align large scans at reduced resolution.

Use several GPUsCopy Anchor Link

devices= shares a scan's views among GPUs in one process. Pass one device or several; jax.devices() passes them all:

import jax

import tomojax as tj

scan = tj.load("scan.nxs")
recon = tj.reconstruct(scan, "fista", nonnegative=True, devices=jax.devices())
result = tj.align(scan, devices=jax.devices())

Each GPU projects its own views and holds the whole volume. Their backprojections are summed on the first GPU, so the result is the one-GPU result up to the order of that sum. Several GPUs therefore add speed, not room for a larger volume, with one exception: in fdk_host and FDK's automatic slabs, each GPU reconstructs its own z-slabs (FDKConfig(devices=...) in fdk_host's FDKHostConfig).

• tj.project, tj.backproject and the fbp (including FDK), cgls and fista methods of tj.reconstruct take devices=. SPDHG-TV runs on one GPU, because each of its steps would sum a whole volume across the GPUs for one block of views, and fbp_host also runs on one.
• Shared projections are held in GPU memory, divided among the GPUs, and are never streamed from the host; asking for both raises ValueError.
• In tj.align, only pose mode's joint pose-and-volume update and cone-beam reconstruction steps are shared. Parallel-beam reconstruction steps, the alternating solver of the other modes and setup calibration run on the first GPU. A checkpoint written on some GPUs resumes on others.
• The GPUs must belong to one process; runs across several machines are not supported.

The tomojax commands have no device option and run on JAX's default device, the first visible GPU (CUDA_VISIBLE_DEVICES chooses which). Call the Python API to use several. On four H100 GPUs, 20 FISTA-TV iterations on the walnut ran 3.3 times faster than on one.