nanodrr

A performance-oriented reimplementation of DiffDRR with the following improvements:

• Optimized, pure PyTorch implementation (~5× faster than DiffDRR at baseline)
• Fused Triton rendering kernel, used by default on CUDA (up to ~80× faster than DiffDRR)
• Modular design (freely swap subjects, extrinsics, and intrinsics during rendering)
• Compatibility with torch.compile and mixed precision
• Extensive type hints with jaxtyping
• Standard Python package structure managed with uv

All projective geometry is implemented internally using the standard Hartley and Zisserman pinhole camera formulation.

Installation

[!NOTE]

On pytorch<2.9, torch.compile with bfloat16 is slower than eager for the pure PyTorch backend due to a CUDA graph capture issue (see Benchmarks). The fused Triton backend is unaffected.

To strictly install the renderer:

pip install nanodrr

To install the optional plotting or 3D visualization module:

pip install "nanodrr[plot]"   # 2D visualization (matplotlib, opencv)
pip install "nanodrr[scene]"  # 3D visualization (VTK, PyVista)
pip install "nanodrr[all]"    # All extras

Benchmarks

[!IMPORTANT]

• ~5× faster than DiffDRR with the pure PyTorch backend (1,093 FPS vs 223 FPS)
• ~57× faster with the fused Triton kernel, without compilation (12,600 FPS vs 223 FPS)
• ~80× faster with torch.compile and bfloat16 (17,820 FPS vs 223 FPS)
• ~3.5× less memory than DiffDRR (322 MB vs 1,170 MB peak reserved with bfloat16 + compile)

FPS computed from per-frame GPU kernel time (torch.profiler); median wall-clock timings are also recorded in tests/benchmark/benchmark.csv. Benchmarked by rendering 200×200 DRRs on an NVIDIA RTX 6000 Ada (48 GB) with Python 3.12. Compile represents torch.compile(mode="reduce-overhead", fullgraph=True). Full experiment at tests/benchmark/.

Docs

To test the docs locally, run

uv run --group docs jupyter nbconvert --to markdown tutorials/*.ipynb --output-dir docs/tutorials/
uv run --group docs zensical serve

Roadmap

• Implement a fully optimized renderer
• Port strictly necessary modules from DiffDRR (e.g., SE(3) utilities, loss functions, and 2D plotting)
• Migrate 3D plotting functions to an optional module
• Integrate with xvr to speed up network training and registration
• Integrate with polypose to speed up registration
• Release as v1.0.0 of DiffDRR!