Code for KiloNeRF: Speeding up Neural Radiance Fields with Thousands of Tiny MLPs

KiloNeRF: Speeding up Neural Radiance Fields with Thousands of Tiny MLPs

Check out the paper on arXiv: https://arxiv.org/abs/2103.13744

KiloNeRF interactive demo

This repo contains the code for KiloNeRF, together with instructions on how to download pretrained models and datasets. Additionally, we provide a viewer for interactive visualization of KiloNeRF scenes. We further improved the implementation and KiloNeRF now runs ~5 times faster than the numbers we report in the first arXiv version of the paper. As a consequence the Lego scene can now be rendered at around 50 FPS.

Prerequisites

  • OS: Ubuntu 20.04.2 LTS
  • GPU: >= NVIDIA GTX 1080 Ti with >= 460.73.01 driver
  • Python package manager conda

Setup

Open a terminal in the root directory of this repo and execute export KILONERF_HOME=$PWD

Install OpenGL and GLUT development files
sudo apt install libgl-dev freeglut3-dev

Install Python packages
conda env create -f $KILONERF_HOME/environment.yml

Activate kilonerf environment
source activate kilonerf

CUDA extension installation

You can either install our pre-compiled CUDA extension or compile the extension yourself. Only compiling it yourself will allow you to make changes to the CUDA code but is more tedious.

Option A: Install pre-compiled CUDA extension

Install pre-compiled CUDA extension
pip install $KILONERF_HOME/cuda/dist/kilonerf_cuda-0.0.0-cp38-cp38-linux_x86_64.whl

Option B: Build CUDA extension yourself

Install CUDA development kit and restart your bash:

wget https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda_11.1.1_455.32.00_linux.run
sudo sh cuda_11.1.1_455.32.00_linux.run
echo -e "\nexport PATH=\"/usr/local/cuda/bin:\$PATH\"" >> ~/.bashrc
echo "export LD_LIBRARY_PATH=\"/usr/local/cuda/lib64:\$LD_LIBRARY_PATH\"" >> ~/.bashrc

Download magma from http://icl.utk.edu/projectsfiles/magma/downloads/magma-2.5.4.tar.gz then build and install to /usr/local/magma

sudo apt install gfortran libopenblas-dev
wget http://icl.utk.edu/projectsfiles/magma/downloads/magma-2.5.4.tar.gz
tar -zxvf magma-2.5.4.tar.gz
cd magma-2.5.4
cp make.inc-examples/make.inc.openblas make.inc
export GPU_TARGET="Maxwell Pascal Volta Turing Ampere"
export CUDADIR=/usr/local/cuda
export OPENBLASDIR="/usr"
make
sudo -E make install prefix=/usr/local/magma

For further information on installing magma see: http://icl.cs.utk.edu/projectsfiles/magma/doxygen/installing.html

Finally compile KiloNeRF's C++/CUDA code

cd $KILONERF_HOME/cuda
python setup.py develop

Download pretrained models

We provide pretrained KiloNeRF models for the following scenes: Synthetic_NeRF_Chair, Synthetic_NeRF_Lego, Synthetic_NeRF_Ship, Synthetic_NSVF_Palace, Synthetic_NSVF_Robot

cd $KILONERF_HOME
mkdir logs
cd logs
wget https://www.dropbox.com/s/eqvf3x23qbubr9p/kilonerf-pretrained.tar.gz?dl=1 --output-document=paper.tar.gz
tar -xf paper.tar.gz

Download NSVF datasets

Credit to NSVF authors for providing their datasets: https://github.com/facebookresearch/NSVF

cd $KILONERF_HOME/data/nsvf
wget https://dl.fbaipublicfiles.com/nsvf/dataset/Synthetic_NSVF.zip && unzip -n Synthetic_NSVF.zip
wget https://dl.fbaipublicfiles.com/nsvf/dataset/Synthetic_NeRF.zip && unzip -n Synthetic_NeRF.zip
wget https://dl.fbaipublicfiles.com/nsvf/dataset/BlendedMVS.zip && unzip -n BlendedMVS.zip
wget https://dl.fbaipublicfiles.com/nsvf/dataset/TanksAndTemple.zip && unzip -n TanksAndTemple.zip

Since we slightly adjusted the bounding boxes for some scenes, it is important that you use the provided unzip argument to avoid overwriting our bounding boxes.

Usage

To benchmark a trained model run:
bash benchmark.sh

You can launch the interactive viewer by running:
bash render_to_screen.sh

To train a model yourself run
bash train.sh

The default dataset is Synthetic_NeRF_Lego, you can adjust the dataset by setting the dataset variable in the respective script.

Owner
Comments
  • How to decide resolution for tiny MLPs in each scene

    How to decide resolution for tiny MLPs in each scene

    Thanks a lot for your great work!

    I noticed the resolution (fixed_resolution) for tiny MLPs defined in cfgs/paper/distill/$SEQ.yaml is not the same across scenes: for lego, [9,16,10], for hotdog: [16,16, 6].

    So, I wonder how did you determine the resolution for each scene.

    Here are my naive approaches: a. just optimize for each scene using grid search b. define the size for a voxel in each scene, then derive the resolution based on global_{max, min} and voxel size

    • e.g. ceil ((global_max - global_min) / voxel size)

    I would appreciate if you have any suggestions.

  • Resolution in pretrain_occupancy cfg file

    Resolution in pretrain_occupancy cfg file

    Hi! First of all, I would like to thank you for your incredible work speeding up NeRF.

    I'm trying to train a new model, but I'm struggling with the configuration. I've read the paper but I can't seem to understand what 'resolution' on cfg files stands for, so I'm not sure which values I should have there for my model. If you could guide me a little with this I would be immensely grateful.

    Thank you.

  • The number of layers and hidden dim are fixed in the CUDA implementation

    The number of layers and hidden dim are fixed in the CUDA implementation

    Hi Christian @creiser !

    Thank you a lot for the great work and the open-sourced code! It is very helpful in the NeRF community!

    I noticed that the number of layers and hidden dims seems to be fixed in the CUDA implementation, regardless of the number of layers and hidden dim we set in the .yaml file.

    If I understand it correctly:

    I was wondering:

    • The reason why you fixed them here (for less pre-allocated memory?)
    • Do you have any plan to extend it to support an arbitrary number of layers and hidden dim?

    Thank you!

  • Inconsistent LPIPS implementation?

    Inconsistent LPIPS implementation?

    Really great work!

    I have a question about current LPIPS implementation. From the lpips document, the input RGB images should be in scale [-1,+1]. But it seems that the current implementation feed RGB in [0, 1] to lpips.

    I have also traced the baseline NSVF's lpips evaluation and they scale the RGB from [0, 1] to [-1, 1] (see their code).

    I would appreciate clarification about this. Thanks!

  • Running our of RAM memory during distillation

    Running our of RAM memory during distillation

    Hi! Firstly, I would like to thank you for work on KiloNeRF and the benefits it provides over basic NeRFs.

    I added another output layer to the base NeRF and now the local_distill.py is always killed during its execution. My system reports that the RAM is running out of memory. Did you experience anything similar and know how to fix that? Is maybe the network that has to be distilled too big?

  • Low PSNR on custom dataset

    Low PSNR on custom dataset

    Hello, first of all thanks for your very interesting work.

    I'm trying to train a very simple custom scene, but the PSNR remains low even after many iterations on the Vanilla NeRF. The Lego data set has a mean PSNR of 16 on the test set after only 500 iterations, while mine has a mean PSNR of 12 after 10k iterations. I generated my dataset with BlenderProc using the same poses of the Lego dataset, and I checked that the intrinsics, bbox and poses are taken correctly.

    These are two images of the train set:

    And these are the results after 10k iterations on the test set:

    And if I try to render images of the train set the results are slightly better, but still not acceptable.

    I can see during the training that the Loss is decreasing and the PSNR is increasing but not as fast as it did with the Lego data set, and they are oscillating a lot.

    Screenshot from 2021-11-15 12-37-18

    Do you have any suggestion to fix this behavior? Thanks a lot!

  • bmm vs cuda implementation

    bmm vs cuda implementation

    Hi, thanks for opening the source code!

    I'm just curious about the acceleration speed of the cuda implementation compared to pytorch.bmm operation, if the input to each MLPs are equal. The test() part code in multi_module.py cannot run successfully due to some flags, and I have no idea how to measure the speed of the cuda implementation against bmm. Could you please give me some guidance? Thanks!

  • CUDA error at /home/chris/anti/cuda/render_to_screen.cpp:113 code=999(cudaErrorUnknown)

    CUDA error at /home/chris/anti/cuda/render_to_screen.cpp:113 code=999(cudaErrorUnknown)

    Hi! I met this CUDA error while running render_to_screen.sh: CUDA error at /home/chris/anti/cuda/render_to_screen.cpp:113 code=999(cudaErrorUnknown) "cudaGraphicsGLRegisterBuffer(&cuda_pbo_resource, pbo, cudaGraphicsMapFlagsWriteDiscard)" render_to_screen.sh: line 3: 28062 Segmentation fault (core dumped) python run_nerf.py cfgs/paper/finetune/$DATASET.yaml -rcfg cfgs/render/render_to_screen.yaml

    I'm running kilonerf on Ubuntu18.04, CUDA11.1, GPU is A6000. Could you please help me with this? Thank you very much!

    Here's the output:

    (kilonerf) [email protected]:~/drivers/5/kilonerf$ bash render_to_screen.sh
    auto log path: logs/paper/finetune/Synthetic_NeRF_Lego
    {'checkpoint_interval': 50000, 'chunk_size': 40000, 'distilled_cfg_path': 'cfgs/paper/distill/Synthetic_NeRF_Lego.yaml', 'distilled_checkpoint_path': 'logs/paper/distill/Synthetic_NeRF_Lego/checkpoint.pth', 'initial_learning_rate': 0.001, 'iterations': 1000000, 'l2_regularization_lambda': 1e-06, 'learing_rate_decay_rate': 500, 'no_batching': True, 'num_rays_per_batch': 8192, 'num_samples_per_ray': 384, 'occupancy_cfg_path': 'cfgs/paper/pretrain_occupancy/Synthetic_NeRF_Lego.yaml', 'occupancy_log_path': 'logs/paper/pretrain_occupancy/Synthetic_NeRF_Lego/occupancy.pth', 'perturb': 1.0, 'precrop_fraction': 0.5, 'precrop_iterations': 0, 'raw_noise_std': 0.0, 'render_only': False, 'no_color_sigmoid': False, 'render_test': True, 'render_factor': 0, 'testskip': 8, 'deepvoxels_shape': 'greek', 'blender_white_background': True, 'blender_half_res': False, 'llff_factor': 8, 'llff_no_ndc': False, 'llff_lindisp': False, 'llff_spherify': False, 'llff_hold': False, 'print_interval': 100, 'render_testset_interval': 10000, 'render_video_interval': 100000000, 'network_chunk_size': 65536, 'rng_seed': 0, 'use_same_initialization_for_all_networks': False, 'use_initialization_fix': False, 'num_importance_samples_per_ray': 0, 'model_type': 'multi_network', 'random_direction_probability': -1, 'von_mises_kappa': -1, 'view_dependent_dropout_probability': -1}
    Using GPU: RTX A6000
    /home/nesc525/drivers/5/kilonerf/utils.py:254: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.
      return np.array([[float(w) for w in line.strip().split()] for line in open(path)]).astype(np.float32)
    Loaded a NSVF-style dataset (138, 800, 800, 4) (138, 4, 4) (0,) data/nsvf/Synthetic_NeRF/Lego
    (100,) (13,) (25,)
    Converting alpha to white.
    global_domain_min: [-0.67 -1.2  -0.37], global_domain_max: [0.67 1.2  1.03], near: 2.0, far: 6.0, background_color: tensor([1., 1., 1.])
    Loading logs/paper/finetune/Synthetic_NeRF_Lego/checkpoint_1000000.pth
    Loading occupancy grid from logs/paper/pretrain_occupancy/Synthetic_NeRF_Lego/occupancy.pth
    CUDA error at /home/chris/anti/cuda/render_to_screen.cpp:113 code=999(cudaErrorUnknown) "cudaGraphicsGLRegisterBuffer(&cuda_pbo_resource, pbo, cudaGraphicsMapFlagsWriteDiscard)" 
    render_to_screen.sh: line 3: 28062 Segmentation fault      (core dumped) python run_nerf.py cfgs/paper/finetune/$DATASET.yaml -rcfg cfgs/render/render_to_screen.yaml
    
  • Question on the activation for density in student NN for distillation

    Question on the activation for density in student NN for distillation

    Thanks a lot for your great work.

    When debugging distillation, I found that sometimes alpha from student NN takes the negative value (but very close to 0). This comes from using different activation for density in teacher NN and student NN:

    • teacher NN: Relu (https://github.com/creiser/kilonerf/blob/master/local_distill.py#L335)
    • student NN: leakey Relu (https://github.com/creiser/kilonerf/blob/master/local_distill.py#L114)

    I guess the difference may be minor, but is there any reason for using leaky Relu for student NN?

  • Best way to train model for LLFF dataset

    Best way to train model for LLFF dataset

    Thank you for your contribution.

    I want to train the KiloNeRF for LLFF dataset.

    How do I set the best configuration for the best training in LLFF dataset? Such as follows: Resolution, ,Global_domain_min, Global_domain_max, fixed_resolution

  • multi-network using for the distill procedure

    multi-network using for the distill procedure

    Hi! Firstly, I would like to thank you for your outstanding work speeding up NeRF.

    I have trained a multi-network in the pretrain procedure(with 27 middle MLP).I want to know can I use this model as the pretrain in your code during the distill procedure.

Open source repository for the code accompanying the paper 'Non-Rigid Neural Radiance Fields Reconstruction and Novel View Synthesis of a Deforming Scene from Monocular Video'.
Open source repository for the code accompanying the paper 'Non-Rigid Neural Radiance Fields Reconstruction and Novel View Synthesis of a Deforming Scene from Monocular Video'.

Non-Rigid Neural Radiance Fields This is the official repository for the project "Non-Rigid Neural Radiance Fields: Reconstruction and Novel View Synt

Sep 22, 2022
Code release for DS-NeRF (Depth-supervised Neural Radiance Fields)
Code release for DS-NeRF (Depth-supervised Neural Radiance Fields)

Depth-supervised NeRF: Fewer Views and Faster Training for Free Project | Paper | YouTube Pytorch implementation of our method for learning neural rad

Sep 25, 2022
This repository contains the source code for the paper "DONeRF: Towards Real-Time Rendering of Compact Neural Radiance Fields using Depth Oracle Networks",
This repository contains the source code for the paper

DONeRF: Towards Real-Time Rendering of Compact Neural Radiance Fields using Depth Oracle Networks Project Page | Video | Presentation | Paper | Data L

Sep 23, 2022
This is the code for "HyperNeRF: A Higher-Dimensional Representation for Topologically Varying Neural Radiance Fields".

HyperNeRF: A Higher-Dimensional Representation for Topologically Varying Neural Radiance Fields This is the code for "HyperNeRF: A Higher-Dimensional

Sep 19, 2022
Code release for NeRF (Neural Radiance Fields)
Code release for NeRF (Neural Radiance Fields)

NeRF: Neural Radiance Fields Project Page | Video | Paper | Data Tensorflow implementation of optimizing a neural representation for a single scene an

Sep 21, 2022
(Arxiv 2021) NeRF--: Neural Radiance Fields Without Known Camera Parameters

NeRF--: Neural Radiance Fields Without Known Camera Parameters Project Page | Arxiv | Colab Notebook | Data Zirui Wang¹, Shangzhe Wu², Weidi Xie², Min

Sep 21, 2022
Unofficial & improved implementation of NeRF--: Neural Radiance Fields Without Known Camera Parameters
Unofficial & improved implementation of NeRF--: Neural Radiance Fields Without Known Camera Parameters

[Unofficial code-base] NeRF--: Neural Radiance Fields Without Known Camera Parameters [ Project | Paper | Official code base ] ⬅️ Thanks the original

Sep 20, 2022
Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields.
Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields.

This repository contains the code release for Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields. This implementation is written in JAX, and is a fork of Google's JaxNeRF implementation. Contact Jon Barron if you encounter any issues.

Sep 26, 2022
This repository contains a PyTorch implementation of "AD-NeRF: Audio Driven Neural Radiance Fields for Talking Head Synthesis".
This repository contains a PyTorch implementation of

AD-NeRF: Audio Driven Neural Radiance Fields for Talking Head Synthesis | Project Page | Paper | PyTorch implementation for the paper "AD-NeRF: Audio

Sep 23, 2022
PyTorch implementation for MINE: Continuous-Depth MPI with Neural Radiance Fields
PyTorch implementation for  MINE: Continuous-Depth MPI with Neural Radiance Fields

MINE: Continuous-Depth MPI with Neural Radiance Fields Project Page | Video PyTorch implementation for our ICCV 2021 paper. MINE: Towards Continuous D

Sep 16, 2022
BARF: Bundle-Adjusting Neural Radiance Fields 🤮 (ICCV 2021 oral)

BARF ?? : Bundle-Adjusting Neural Radiance Fields Chen-Hsuan Lin, Wei-Chiu Ma, Antonio Torralba, and Simon Lucey IEEE International Conference on Comp

Sep 17, 2022
[ICCV21] Self-Calibrating Neural Radiance Fields
[ICCV21] Self-Calibrating Neural Radiance Fields

Self-Calibrating Neural Radiance Fields, ICCV, 2021 Project Page | Paper | Video Author Information Yoonwoo Jeong [Google Scholar] Seokjun Ahn [Google

Sep 22, 2022
[ICCV 2021 Oral] NerfingMVS: Guided Optimization of Neural Radiance Fields for Indoor Multi-view Stereo
[ICCV 2021 Oral] NerfingMVS: Guided Optimization of Neural Radiance Fields for Indoor Multi-view Stereo

NerfingMVS Project Page | Paper | Video | Data NerfingMVS: Guided Optimization of Neural Radiance Fields for Indoor Multi-view Stereo Yi Wei, Shaohui

Sep 27, 2022
A PyTorch implementation of NeRF (Neural Radiance Fields) that reproduces the results.
A PyTorch implementation of NeRF (Neural Radiance Fields) that reproduces the results.

NeRF-pytorch NeRF (Neural Radiance Fields) is a method that achieves state-of-the-art results for synthesizing novel views of complex scenes. Here are

Sep 22, 2022
D-NeRF: Neural Radiance Fields for Dynamic Scenes
 D-NeRF: Neural Radiance Fields for Dynamic Scenes

D-NeRF: Neural Radiance Fields for Dynamic Scenes [Project] [Paper] D-NeRF is a method for synthesizing novel views, at an arbitrary point in time, of

Sep 7, 2022
pixelNeRF: Neural Radiance Fields from One or Few Images
 pixelNeRF: Neural Radiance Fields from One or Few Images

pixelNeRF: Neural Radiance Fields from One or Few Images Alex Yu, Vickie Ye, Matthew Tancik, Angjoo Kanazawa UC Berkeley arXiv: http://arxiv.org/abs/2

Sep 18, 2022
A PyTorch re-implementation of Neural Radiance Fields
A PyTorch re-implementation of Neural Radiance Fields

nerf-pytorch A PyTorch re-implementation Project | Video | Paper NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis Ben Mildenhall

Sep 17, 2022
[ICCV'21] UNISURF: Unifying Neural Implicit Surfaces and Radiance Fields for Multi-View Reconstruction
[ICCV'21] UNISURF: Unifying Neural Implicit Surfaces and Radiance Fields for Multi-View Reconstruction

UNISURF: Unifying Neural Implicit Surfaces and Radiance Fields for Multi-View Reconstruction Project Page | Paper | Supplementary | Video This reposit

Sep 16, 2022
This is a JAX implementation of Neural Radiance Fields for learning purposes.

learn-nerf This is a JAX implementation of Neural Radiance Fields for learning purposes. I've been curious about NeRF and its follow-up work for a whi

Aug 22, 2022