0.5.0: Rewrite all the underlying CUDA. Speedup and Benchmarking. (#182)

* importance_sampling with test

* package importance_sampling

* compute_intervals tested and packaged

* compute_intervals_v2

* bicycle is failing

* fix cut in compute_intervals_v2, test pass for rendering

* hacky way to get opaque_bkgd work

* reorg ING

* PackedRaySegmentsSpec

* chunk_ids -> ray_ids

* binary -> occupied

* test_traverse_grid_basic checked

* fix traverse_grid with step size, checked

* support max_step_size, not verified

* _cuda and cuda; upgrade ray_marching

* inclusive scan

* test_exclusive_sum but seems to have numeric error

* inclusive_sum_backward verified

* exclusive sum backward

* merge fwd and bwd for scan

* inclusive & exclusive prod verified

* support normal scan with torch funcs

* rendering and tests

* a bit clean up

* importance_sampling verified

* stratified for importance_sampling

* importance_sampling in pdf.py

* RaySegmentsSpec in data_specs; fix various bugs

* verified with _proposal_packed.py

* importance sampling support batch input/output. need to verify

* prop script with batch samples

* try to use cumsum  instead of cumprod

* searchsorted

* benchmarking prop

* ray_aabb_intersect untested

* update prop benchmark numbers

* minor fixes

* batched ray_aabb_intersect

* ray_aabb_intersect and traverse with grid(s)

* tiny optimize for traverse_grids kernels

* traverse_grids return intervals and samples

* cub not verified

* cleanup

* propnet and occgrid as estimators

* training print iters 10k

* prop is good now

* benchmark in google sheet.

* really cleanup: scan.py and test

* pack.py and test

* rendering and test

* data_specs.py and pdf.py docs

* data_specs.py and pdf.py docs

* init and headers

* grid.py and test for it

* occ grid docs

* generated docs

* example docs for pack and scan function.

* doc fix for volrend.py

* doc fix for pdf.py

* fix doc for rendering function

* docs

* propnet docs

* update scripts

* docs: index.rst

* methodology docs

* docs for examples

* mlp nerf script

* update t-nerf script

* rename dnerf to tnerf

* misc update

* bug fix: pdf_loss with test

* minor fix

* update readme with submodules

* fix format

* update gitingore file

* fix doc failure. teaser png to jpg

* docs in examples/

.gitignore

@@ -117,6 +117,11 @@ venv.bak/
 # vscode
 .vsocde
 
-benchmarks/
 outputs/
-data
\ No newline at end of file
+data
+
+benchmarks/*
+!benchmarks/barf
+!benchmarks/kplanes
+!benchmarks/tensorf
+!benchmarks/tineuvox
\ No newline at end of file

.gitmodules

 [submodule "examples/pycolmap"]
 	path = examples/pycolmap
-	url = https://github.com/rmbrualla/pycolmap.git
\ No newline at end of file
+	url = https://github.com/rmbrualla/pycolmap.git
+
+[submodule "benchmarks/tensorf"]
+	path = benchmarks/tensorf
+	url = https://github.com/liruilong940607/tensorf.git
+	branch = nerfacc
+
+[submodule "benchmarks/barf"]
+	path = benchmarks/barf
+	url = https://github.com/liruilong940607/barf.git
+	branch = nerfacc
+
+[submodule "benchmarks/tineuvox"]
+	path = benchmarks/tineuvox
+	url = https://github.com/liruilong940607/tineuvox.git
+	branch = nerfacc
+
+[submodule "benchmarks/kplanes"]
+	path = benchmarks/kplanes
+	url = https://github.com/liruilong940607/kplanes.git
+	branch = nerfacc

MANIFEST.in

 include nerfacc/cuda/csrc/include/*
 include nerfacc/cuda/csrc/*
+
+include nerfacc/_cuda/csrc/include/*
+include nerfacc/_cuda/csrc/*

README.md

@@ -16,19 +16,14 @@
 
 https://www.nerfacc.com/
 
-NerfAcc is a PyTorch Nerf acceleration toolbox for both training and inference. It focuses on efficient volumetric rendering of radiance fields, which is universal and plug-and-play for most of the NeRFs.
+NerfAcc is a PyTorch Nerf acceleration toolbox for both training and inference. It focus on
+efficient sampling in the volumetric rendering pipeline of radiance fields, which is 
+universal and plug-and-play for most of the NeRFs.
+With minimal modifications to the existing codebases, Nerfacc provides significant speedups 
+in training various recent NeRF papers.
+**And it is pure Python interface with flexible APIs!**
 
-Using NerfAcc, 
-
-- The `vanilla NeRF` model with 8-layer MLPs can be trained to *better quality* (+~0.5 PNSR)
-  in *1 hour* rather than *days* as in the paper.
-- The `Instant-NGP NeRF` model can be trained to *equal quality* in *4.5 minutes*,
-  comparing to the official pure-CUDA implementation.
-- The `D-NeRF` model for *dynamic* objects can also be trained in *1 hour*
-  rather than *2 days* as in the paper, and with *better quality* (+~2.5 PSNR).
-- Both *bounded* and *unbounded* scenes are supported.
-
-**And it is a pure Python interface with flexible APIs!**
+![Teaser](/docs/source/_static/images/teaser.jpg?raw=true)
 
 ## Installation
 
@@ -62,11 +57,17 @@ pip install nerfacc -f https://nerfacc-bucket.s3.us-west-2.amazonaws.com/whl/tor
 
 ## Usage
 
-The idea of NerfAcc is to perform efficient ray marching and volumetric rendering. So NerfAcc can work with any user-defined radiance field. To plug the NerfAcc rendering pipeline into your code and enjoy the acceleration, you only need to define two functions with your radiance field.
-- `sigma_fn`: Compute density at each sample. It will be used by `nerfacc.ray_marching()` to skip the empty and occluded space during ray marching, which is where the major speedup comes from. 
-- `rgb_sigma_fn`: Compute color and density at each sample. It will be used by `nerfacc.rendering()` to conduct differentiable volumetric rendering. This function will receive gradients to update your network.
+The idea of NerfAcc is to perform efficient volumetric sampling with a computationally cheap estimator to discover surfaces.
+So NerfAcc can work with any user-defined radiance field. To plug the NerfAcc rendering pipeline into your code and enjoy 
+the acceleration, you only need to define two functions with your radience field.
+
+- `sigma_fn`: Compute density at each sample. It will be used by the estimator
+  (e.g., `nerfacc.OccGridEstimator`, `nerfacc.PropNetEstimator`) to discover surfaces. 
+- `rgb_sigma_fn`: Compute color and density at each sample. It will be used by 
+  `nerfacc.rendering` to conduct differentiable volumetric rendering. This function 
+  will receive gradients to update your radiance field.
 
-A simple example is like this:
+An simple example is like this:
 
 ``` python
 import torch
@@ -76,50 +77,39 @@ import nerfacc
 radiance_field = ...  # network: a NeRF model
 rays_o: Tensor = ...  # ray origins. (n_rays, 3)
 rays_d: Tensor = ...  # ray normalized directions. (n_rays, 3)
-optimizer = ...  # optimizer
+optimizer = ...       # optimizer
+
+estimator = nerfacc.OccGridEstimator(...)
 
 def sigma_fn(
     t_starts: Tensor, t_ends:Tensor, ray_indices: Tensor
 ) -> Tensor:
-    """ Query density values from a user-defined radiance field.
-    :params t_starts: Start of the sample interval along the ray. (n_samples, 1).
-    :params t_ends: End of the sample interval along the ray. (n_samples, 1).
-    :params ray_indices: Ray indices that each sample belongs to. (n_samples,).
-    :returns The post-activation density values. (n_samples, 1).
-    """
+    """ Define how to query density for the estimator."""
     t_origins = rays_o[ray_indices]  # (n_samples, 3)
     t_dirs = rays_d[ray_indices]  # (n_samples, 3)
-    positions = t_origins + t_dirs * (t_starts + t_ends) / 2.0
+    positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
     sigmas = radiance_field.query_density(positions) 
-    return sigmas  # (n_samples, 1)
+    return sigmas  # (n_samples,)
 
 def rgb_sigma_fn(
     t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor
 ) -> Tuple[Tensor, Tensor]:
-    """ Query rgb and density values from a user-defined radiance field.
-    :params t_starts: Start of the sample interval along the ray. (n_samples, 1).
-    :params t_ends: End of the sample interval along the ray. (n_samples, 1).
-    :params ray_indices: Ray indices that each sample belongs to. (n_samples,).
-    :returns The post-activation rgb and density values. 
-        (n_samples, 3), (n_samples, 1).
-    """
+    """ Query rgb and density values from a user-defined radiance field. """
     t_origins = rays_o[ray_indices]  # (n_samples, 3)
     t_dirs = rays_d[ray_indices]  # (n_samples, 3)
-    positions = t_origins + t_dirs * (t_starts + t_ends) / 2.0
+    positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
     rgbs, sigmas = radiance_field(positions, condition=t_dirs)  
-    return rgbs, sigmas  # (n_samples, 3), (n_samples, 1)
+    return rgbs, sigmas  # (n_samples, 3), (n_samples,)
 
-# Efficient Raymarching: Skip empty and occluded space, pack samples from all rays.
-# ray_indices: (n_samples,). t_starts: (n_samples, 1). t_ends: (n_samples, 1).
-with torch.no_grad():
-    ray_indices, t_starts, t_ends = nerfacc.ray_marching(
-        rays_o, rays_d, sigma_fn=sigma_fn, near_plane=0.2, far_plane=1.0, 
-        early_stop_eps=1e-4, alpha_thre=1e-2, 
-    )
+# Efficient Raymarching:
+# ray_indices: (n_samples,). t_starts: (n_samples,). t_ends: (n_samples,).
+ray_indices, t_starts, t_ends = estimator.sampling(
+    rays_o, rays_d, sigma_fn=sigma_fn, near_plane=0.2, far_plane=1.0, early_stop_eps=1e-4, alpha_thre=1e-2, 
+)
 
 # Differentiable Volumetric Rendering.
 # colors: (n_rays, 3). opaicity: (n_rays, 1). depth: (n_rays, 1).
-color, opacity, depth = nerfacc.rendering(
+color, opacity, depth, extras = nerfacc.rendering(
     t_starts, t_ends, ray_indices, n_rays=rays_o.shape[0], rgb_sigma_fn=rgb_sigma_fn
 )
 
@@ -132,6 +122,8 @@ optimizer.step()
 
 ## Examples: 
 
+See full benchmarking here: https://www.nerfacc.com/en/latest/examples/
+
 Before running those example scripts, please check the script about which dataset is needed, and download the dataset first. You could use `--data_root` to specify the path.
 
 ```bash
@@ -139,38 +131,80 @@ Before running those example scripts, please check the script about which datase
 git clone --recursive git://github.com/KAIR-BAIR/nerfacc/
 ```
 
+### Static NeRFs
+
+Instant-NGP on NeRF-Synthetic dataset with better performance in 4.5 minutes.
 ``` bash
-# Instant-NGP NeRF in 4.5 minutes with reproduced performance!
-# See results at here: https://www.nerfacc.com/en/latest/examples/ngp.html
-python examples/train_ngp_nerf.py --train_split train --scene lego
+# Occupancy Grid Estimator
+python examples/train_ngp_nerf_occ.py --scene lego --data_root data/nerf_synthetic
+# Proposal Net Estimator
+python examples/train_ngp_nerf_prop.py --scene lego --data_root data/nerf_synthetic
 ```
 
+Instant-NGP on Mip-NeRF 360 dataset with better performance in 5 minutes.
 ``` bash
-# Vanilla MLP NeRF in 1 hour with better performance!
-# See results at here: https://www.nerfacc.com/en/latest/examples/vanilla.html
-python examples/train_mlp_nerf.py --train_split train --scene lego
+# Occupancy Grid Estimator
+python examples/train_ngp_nerf_occ.py --scene garden --data_root data/360_v2
+# Proposal Net Estimator
+python examples/train_ngp_nerf_prop.py --scene garden --data_root data/360_v2
 ```
 
+Vanilla MLP NeRF on NeRF-Synthetic dataset in an hour.
+``` bash
+# Occupancy Grid Estimator
+python examples/train_mlp_nerf.py --scene lego --data_root data/nerf_synthetic
+```
+
+TensoRF on Tanks&Temple and NeRF-Synthetic datasets (plugin in the official codebase).
+``` bash
+cd benchmarks/tensorf/
+# (set up the environment for that repo)
+bash script.sh nerfsyn-nerfacc-occgrid 0
+bash script.sh tt-nerfacc-occgrid 0
+```
+
+### Dynamic NeRFs
+T-NeRF on D-NeRF dataset in an hour.
+``` bash
+# Occupancy Grid Estimator
+python examples/train_mlp_tnerf.py --scene lego --data_root data/dnerf
+```
+
+K-Planes on D-NeRF dataset (plugin in the official codebase).
 ```bash
-# D-NeRF for Dynamic objects in 1 hour with better performance!
-# See results at here: https://www.nerfacc.com/en/latest/examples/dnerf.html
-python examples/train_mlp_dnerf.py --train_split train --scene lego
+cd benchmarks/kplanes/
+# (set up the environment for that repo)
+bash script.sh dnerf-nerfacc-occgrid 0
 ```
 
+TiNeuVox on HyperNeRF and D-NeRF datasets (plugin in the official codebase).
 ```bash
-# Instant-NGP on unbounded scenes in 20 minutes!
-# See results at here: https://www.nerfacc.com/en/latest/examples/unbounded.html
-python examples/train_ngp_nerf.py --train_split train --scene garden --auto_aabb --unbounded --cone_angle=0.004
+cd benchmarks/tineuvox/
+# (set up the environment for that repo)
+bash script.sh dnerf-nerfacc-occgrid 0
+bash script.sh hypernerf-nerfacc-occgrid 0
+bash script.sh hypernerf-nerfacc-propnet 0
 ```
 
-Used by:
+### Camera Optimization NeRFs
+
+BARF on the NeRF-Synthetic dataset (plugin in the official codebase).
+```bash
+cd benchmarks/barf/
+# (set up the environment for that repo)
+bash script.sh nerfsyn-nerfacc-occgrid 0
+```
+
+### 3rd-Party Usages:
 - [nerfstudio](https://github.com/nerfstudio-project/nerfstudio): A collaboration friendly studio for NeRFs.
+- [sdfstudio](https://autonomousvision.github.io/sdfstudio/): A unified framework for surface reconstruction..
 - [instant-nsr-pl](https://github.com/bennyguo/instant-nsr-pl): NeuS in 10 minutes.
+- [modelscope](https://github.com/modelscope/modelscope/blob/master/modelscope/models/cv/nerf_recon_acc/network/nerf.py): A collection of deep-learning algorithms.
+- [Representing Volumetric Videos as Dynamic MLP Maps, CVPR 2023](https://github.com/zju3dv/mlp_maps)
 
 
 ## Common Installation Issues
 
-
 <details>
     <summary>ImportError: .../csrc.so: undefined symbol</summary>
     If you are installing a pre-built wheel, make sure the Pytorch and CUDA version matchs with the nerfacc version (nerfacc.__version__).
@@ -179,10 +213,10 @@ Used by:
 ## Citation
 
 ```bibtex
-@article{li2022nerfacc,
-  title={NerfAcc: A General NeRF Accleration Toolbox.},
-  author={Li, Ruilong and Tancik, Matthew and Kanazawa, Angjoo},
-  journal={arXiv preprint arXiv:2210.04847},
-  year={2022}
+@article{li2023nerfacc,
+  title={NerfAcc: Efficient Sampling Accelerates NeRFs.},
+  author={Li, Ruilong and Hang Gao and Tancik, Matthew and Kanazawa, Angjoo},
+  journal={TBD},
+  year={2023}
 }
 ```

barf @ 90440d97

+Subproject commit 90440d975fc76b3559126992b2fbce27dd02456f

kplanes @ b97bc2ee

+Subproject commit b97bc2eefc18f00cd54833800e7fc1072e58be51

tensorf @ f2d35087

+Subproject commit f2d350873c54f249e64b6e745919b6a94bf54f1d

tineuvox @ 09998587

+Subproject commit 0999858745577ff32e5226c51c5c78b8315546c8

docs/source/apis/grid.rst → docs/source/apis/estimators.rst

-.. _`Occupancy Grid`:
-
-Occupancy Grid
+.. _`Estimators`:
+ 
+Estimators
 ===================================
 
 .. currentmodule:: nerfacc
 
-.. autoclass:: ContractionType
-    :members:
-
-.. autoclass:: Grid
+.. autoclass:: OccGridEstimator
     :members:
 
-.. autoclass:: OccupancyGrid
+.. autoclass:: PropNetEstimator
     :members:
-

docs/source/apis/generated/nerfacc.exclusive_prod.rst

+nerfacc.exclusive\_prod
+=======================
+
+.. currentmodule:: nerfacc
+
+.. autofunction:: exclusive_prod
\ No newline at end of file

docs/source/apis/generated/nerfacc.exclusive_sum.rst

+nerfacc.exclusive\_sum
+======================
+
+.. currentmodule:: nerfacc
+
+.. autofunction:: exclusive_sum
\ No newline at end of file

docs/source/apis/generated/nerfacc.importance_sampling.rst

+nerfacc.importance\_sampling
+============================
+
+.. currentmodule:: nerfacc
+
+.. autofunction:: importance_sampling
\ No newline at end of file

docs/source/apis/generated/nerfacc.inclusive_prod.rst

+nerfacc.inclusive\_prod
+=======================
+
+.. currentmodule:: nerfacc
+
+.. autofunction:: inclusive_prod
\ No newline at end of file

docs/source/apis/generated/nerfacc.inclusive_sum.rst

+nerfacc.inclusive\_sum
+======================
+
+.. currentmodule:: nerfacc
+
+.. autofunction:: inclusive_sum
\ No newline at end of file

docs/source/apis/generated/nerfacc.pack_data.rst → docs/source/apis/generated/nerfacc.pack_info.rst

-nerfacc.pack\_data
+nerfacc.pack\_info
 ==================
 
 .. currentmodule:: nerfacc
 
-.. autofunction:: pack_data
\ No newline at end of file
+.. autofunction:: pack_info
\ No newline at end of file