Nerf.jl a real-time Neural 3D Scene Reconstruction in pure Julia

We demonstrate a set of pure Julia packages to interactively reconstruct and render Neural Radiance Fields (NeRFs), a 3D neural representation of a scene from pictures, using heterogeneous kernel programming. This system was developed at AMD and supports other GPU backends as well.

Representing scenes as Neural Radiance Fields for novel view synthesis has become one of the most popular approaches in deep-learning in contrast to more classical approaches, like triangle-model-based scene reconstruction. However, training and evaluation of Neural Radiance Fields can be costly and slow, which limits their real-time applicability.

Our package implements many state-of-the-art NeRF algorithms, including Instant Neural Graphics Primitives with Multiresolution Hash Encoding. While a similar open-source C++ and CUDA system exists, we implemented this version from scratch in pure Julia, using heterogeneous kernel programming to support many backends.

We implement efficient kernels for ray marching, including acceleration structures, multiresolution hash encoding and alpha compositing and provide gradient kernels for them. Additionally, we integrate those kernels with Julia's AD ecosystem using ChainRules.jl. This allows them to be used independently and with the AD system of your choice.

For visualization and interaction, we bundle all this into an interactive OpenGL application that enables real-time tweaking of training parameters.

Additional features include the ability to record videos from custom camera paths, convert a trained NeRF to a triangle mesh using Marching Cubes and Marching Tetrahedra algorithms.

During the development of this package, we also improved the Julia GPU stack. We now support ROCm artifacts for AMDGPU.jl and improved the support and stability for RDNA2 GPUs.

In conclusion, we provide a more user-friendly experience, running both on AMDGPU.jl and CUDA.jl. We believe this code is much simpler than its C++/CUDA counterpart, and still delivers competitive performance. We hope this can serve as a good example of how to integrate high-performance 3D graphics and Machine Learning in a GPU-agnostic manner, and contribute to the expansion of the Julia ecosystem.