A PyTorch re-implementation of Neural Radiance Fields.
The current implementation is blazingly fast! (Thorough benchmark to come, but ~2x faster)
The NeRF code release has an accompanying Colab notebook, that showcases training a feature-limited version of NeRF on a "tiny" scene. It's equivalent PyTorch notebook can be found at the following URL:
https://colab.research.google.com/drive/1rO8xo0TemN67d4mTpakrKrLp03b9bgCX
To train a "full" NeRF model (i.e., using 3D coordinates as well as ray directions, and the hierarchical sampling procedure), first setup dependencies. In a new conda or virtualenv environment, run
pip install requirements.txt
Importantly, install torchsearchsorted by following instructions from their README.
Once everything is setup, to run experiments, first edit config/lego.yml to specify your own parameters.
The training script can be invoked by running
python train_nerf.py --config config/lego.yml
Optionally, if resuming training from a previous checkpoint, run
python train_nerf.py --config config/lego.yml --load-checkpoint path/to/checkpoint.ckpt
A Colab notebook for the full NeRF model (albeit on low-resolution data) can be accessed here.
All said, this is not an official code release, and is instead a reproduction from the original code (released by the authors here).
The code is thoroughly tested (to the best of my abilities) to match the original implementation (and be much faster)! In particular, I have ensured that
- Every individual module exactly (numerically) matches that of the TensorFlow implementation. This Colab notebook has all the tests, matching op for op (but is very scratchy to look at)!
- Training works as expected (for Lego and LLFF scenes).
The organization of code WILL change around a lot, because I'm actively experimenting with this.
Pretrained models: I am running a few large-scale experiments, and I hope to release models sometime in the end of April.
Feel free to raise GitHub issues if you find anything concerning. Pull requests adding additional features are welcome too.
nerf-pytorch is available under the MIT License. For more details see: LICENSE and ACKNOWLEDGEMENTS.
Also, a shoutout to yenchenlin for his cool PyTorch implementation, whose volume rendering function replaced mine (my initial impl was inefficient in comparison).