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 (~40× faster than DiffDRR)
• Fused Metal shader (experimental), used by default on MPS (~4× faster than PyTorch)
• Modular design (freely swap subjects, extrinsics, and intrinsics during rendering)
• Compatibility with torch.compile and mixed precision
• Extensive type hints with jaxtyping
• Runs on CUDA, Apple silicon (MPS), and CPU
• Standard Python package structure managed with uv

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

Installation

Note

MPS-accelerated differentiable rendering requires pytorch>=2.13, which is the first release with native MPS kernels for grid_sample backwards and 3D nearest-neighbor sampling.

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,048 FPS vs 226 FPS)
• ~39× faster with the fused Triton kernel, without compilation (8,717 FPS vs 226 FPS)
• ~40× faster with torch.compile and bfloat16 (9,008 FPS vs 226 FPS)
• ~4× less memory than DiffDRR (276 MB vs 1,092 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!