FlashDRR

Fast differentiable DRR raycasting and volume rendering for PyTorch.

Full API reference, and guides available at https://pcarnah.github.io/flashdrr/.

FlashDRR provides fast, differentiable raycasters that turn CT volumes into synthesized X‑ray projections (Digitally Reconstructed Radiographs, DRR) while keeping gradients flowing through both the volume and the camera parameters. It is designed for deep-learning workloads such as 2D/3D registration, pose estimation, and synthetic‑image generation from medical volumes.

Features

• Differentiable DRR raycasting — flashdrr.rendering.VolumeRaycaster renders physically‑based DRRs using the Beer‑Lambert law with fully differentiable ray sampling.
• Fast Triton‑fused renderer — flashdrr.rendering.FusedVolumeRenderer runs a CUDA kernel for near‑instant, backprop‑free (or autograd-aware) projection at high resolution.
• C‑arm / camera helpers — sample realistic C‑arm views (carm_to_camera_params, get_random_carm_views) and build VTK‑compatible camera matrices in RAS space (get_vtk_view_mat).
• Transfer functions — map CT Hounsfield units to density/attenuation via flashdrr.rendering.piecewise_linear_channelwise.
• Lightweight volume utilities — flashdrr.utils provides dimension helpers (make_nd), HU normalization, and voxel‑scale handling.

Benchmark

The following trends summarize head-to-head measurements of VolumeRaycaster (PyTorch grid_sample based) against the Triton-fused FusedVolumeRenderer across resolutions, ray-sample counts, and volume sizes. All tests performed using an RTX 5070Ti 16GB.

Key Performance Trends

1. Memory Scaling: Linear Allocation vs. Near-Constant Footprint

   • Without Triton: Peak VRAM scales directly with total samples ($\text{Rays} \times \text{Samples/Ray}$). At $2048 \times 2048$ resolution with 256 ray samples, peak memory hits 24.8 GB in evaluation mode, leading to out-of-memory (OOM) failures at $\geq 384$ samples.
   • With Triton: Fused rendering kernels avoid materializing intermediate ray-sample tensors in HBM. Memory footprint remains below 430 MiB across all tested configurations — a 99.5% memory reduction at high workloads.
   • Volume Size Impact: Scaling volume grid resolution from $128^3$ to $256^3$ adds a fixed baseline memory overhead (~100–300 MiB) for tensor storage, but Triton's intermediate sample memory remains invariant to volume size.

2. Time Scaling: Sub-Linear Latency vs. Memory-Wall Thrashing

   • Extreme Acceleration: At $2048 \times 2048$ resolution (256 samples, eval), Triton reduces iteration time from 1,804.32 ms to 9.16 ms — a $197\times$ speedup.
   • Non-Linear Degradation without Triton: Beyond $1024 \times 1024$, standard PyTorch execution time degrades super-linearly due to memory bandwidth limits and cache thrashing. Triton maintains predictable linear scaling proportional to total rays rendered.

3. Train vs. Eval Dynamics

   • Memory Overhead: Without Triton, evaluation and training both consume massive memory ($\approx 12.4$ GB at $1024\times 1024$ with 512 samples). With Triton, training memory increases only slightly over evaluation (121.6 MiB train vs 110.7 MiB eval) due to lightweight gradient buffer retention.
     • Backward Pass Latency: Training mode adds a larger relative compute penalty in Triton than in standard PyTorch (~$3.05\times$ overhead vs ~$1.31\times$ at $1024\times 1024$, 512 samples). This stems from executing backward gradient kernels for custom operators, though Triton train step time (15.4 ms) remains dramatically faster than non-Triton eval time (82.0 ms).

Configuration snapshot (1024×1024, 512 samples)

Installation Instructions

The project is available via PyPI.

pip install flashdrr

To get the latest master:

pip install git+https://github.com/pcarnah/flashdrr.git@master#egg=flashdrr

CUDA builds

PyTorch ships separate wheels per CUDA version. By default pip install flashdrr pulls the CPU build of torch. To select a CUDA build, install the corresponding extra and point pip at the matching PyTorch index, for example:

pip install flashdrr[cu128] --extra-index-url https://download.pytorch.org/whl/cu128
pip install flashdrr[cu130] --extra-index-url https://download.pytorch.org/whl/cu130

The Triton-fused renderer requires a working Triton installation (the triton-windows fork is used automatically on Windows).

Using uv

uv is supported. CUDA build selection uses the same cu128 / cu130 / cpu extras as with pip — pick one and the matching PyTorch index is selected automatically via [tool.uv.sources]:

uv sync --extra cu128
uv sync --extra cu130
uv sync --extra cpu

Dev-only tooling (e.g. pytest) lives in PEP 735 dependency groups, which must not share names with the CUDA extras — add them with --group:

uv sync --extra cu130 --group test

The optional data extra (NIfTI file support) is just another extra and combines with a CUDA extra:

uv sync --extra cu128 --extra data

Quick example

The snippet below mirrors the end-to-end __main__ in flashdrr/rendering/raycast.py: it loads a CT volume with MONAI, derives the RAS↔IJK affine, builds a C-arm view matrix with get_vtk_view_mat, and renders a DRR with VolumeRaycaster. Two attenuation pipelines are shown from the same loaded volume — a hand-made 5-keypoint transfer function mapping HU → attenuation coefficients via piecewise_linear_channelwise, and a simple linear HU → μ map.

import torch
from monai.transforms import Compose, LoadImage, EnsureChannelFirst, Spacing, ScaleIntensityRange, EnsureType

import flashdrr.rendering as R

# ---------- 1. Load CT volume and build the RAS<->IJK affine ----------
load_tf = Compose([
    LoadImage(),
    EnsureChannelFirst(),
    Spacing([2.5, 2.5, 3.0]),
    ScaleIntensityRange(-3024, 3024, 0, 1, clip=True),
    EnsureType(),
])
vol = load_tf('CTChest.nii.gz').unsqueeze(0)                    # (1, 1, D, H, W), MetaTensor in HU
ijk2ras = vol.meta['affine']                                    # (4, 4) IJK -> RAS
ras2ijk = torch.linalg.inv(ijk2ras)
vol = vol.cuda()

# ---------- 2. Build a C-arm view matrix in RAS ----------
center_ijk = torch.ones(4, dtype=torch.float64)
center_ijk[:3] = torch.as_tensor(vol.shape[2:]) // 2
center_ras = (ijk2ras @ center_ijk)[:3].float().cuda()
view_mat = R.get_vtk_view_mat(
    cam_pos=(0.0, 1000.0, -130.0),                              # source 1 m in front of the volume
    cam_focal=tuple(center_ras.tolist()),
    cam_viewup=(0.0, 0.0, 1.0),
    device='cuda',
).unsqueeze(0)                                                 # (1, 4, 4)

raycaster = R.VolumeRaycaster(ray_samples=384, resolution=(1024, 1024)).cuda().eval()

# ---------- 3a. Hand-made HU -> attenuation transfer function ----------
# Each keypoint is (HU, attenuation); below -200 HU is air (mu=0),
# soft tissue ramps gently, bone saturates to high attenuation.
tf = torch.tensor([
    [-3500, 0.00],   # air
    [ -200, 0.00],   # lung / soft-tissue boundary
    [  200, 0.05],   # soft tissue
    [ 1535, 0.50],   # cortical bone
    [ 3071, 0.65],   # dense bone / metal
]).cuda()
tf[:, 0] = (tf[:, 0] + 3500) / 7000                            # normalize HU to [0, 1]
xp = tf[:, 0].unsqueeze(0)                                     # (1, K)
yp = tf[:, -1].unsqueeze(0)                                    # (1, K)
mu_tf = R.piecewise_linear_channelwise(vol, xp, yp)            # (1, 1, D, H, W)

drr_tf = raycaster(mu_tf, view_mat=view_mat, ras2ijk=ras2ijk.float().cuda())

# ---------- 3b. Linear HU -> mu map ----------
hu = vol * (3024 - (-3524)) + (-3524)
mu_lin = torch.clamp(0.05 * (1.0 + hu / 800.0), min=0.0)       # (1, 1, D, H, W)

drr_lin = raycaster(mu_lin, view_mat=view_mat, ras2ijk=ras2ijk.float().cuda())

Pass triton=True to raycaster(...) to switch to the fused CUDA/Triton kernel.

Optional dependencies

• nibabel — required only for loading Nifti format data for tests and example scripts (flashdrr[data]).

Attribution

If you use this software in your research, please cite it using the "Cite this repository" button or the provided CITATION.cff.