popcorn-attack / pytorch-nerf

Minimal PyTorch implementations of NeRF and pixelNeRF.

PyTorch NeRF and pixelNeRF

NeRF:

Tiny NeRF:

pixelNeRF:

This repository contains minimal PyTorch implementations of the NeRF model described in "NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis" and the pixelNeRF model described in "pixelNeRF: Neural Radiance Fields from One or Few Images". While there are other PyTorch implementations out there (e.g., this one and this one for NeRF, and the authors' official implementation for pixelNeRF), I personally found them somewhat difficult to follow, so I decided to do a complete rewrite of NeRF myself. I tried to stay as close to the authors' text as possible, and I added comments in the code referring back to the relevant sections/equations in the paper. The final result is a tight 355 lines of heavily commented code (301 sloc—"source lines of code"—on GitHub) all contained in a single file. For comparison, this PyTorch implementation has approximately 970 sloc spread across several files, while this PyTorch implementation has approximately 905 sloc.

run_tiny_nerf.py trains a simplified NeRF model inspired by the "Tiny NeRF" example provided by the NeRF authors. This NeRF model does not use fine sampling and the MLP is smaller, but the code is otherwise identical to the full model code. At only 153 sloc, it might be a good place to start for people who are completely new to NeRF. If you prefer your code more object-oriented, check out run_nerf_alt.py and run_tiny_nerf_alt.py.

A Colab notebook for the full model can be found here, while a notebook for the tiny model can be found here. The generate_nerf_dataset.py script was used to generate the training data of the ShapeNet car (see "Generating the ShapeNet datasets" for additional details).

For the following test view:

run_nerf.py generated the following after 20,100 iterations (a few hours on a P100 GPU):

Loss: 0.00022201683896128088

while run_tiny_nerf.py generated the following after 19,600 iterations (~35 minutes on a P100 GPU):

Loss: 0.0004151524917688221

The advantages of streamlining NeRF's code become readily apparent when trying to extend NeRF. For example, training a pixelNeRF model only required making a few changes to run_nerf.py bringing it to 368 sloc (notebook here). For comparison, the official pixelNeRF implementation has approximately 1,300 pixelNeRF-specific (i.e., not related to the image encoder or dataset) sloc spread across several files. The generate_pixelnerf_dataset.py script was used to generate the training data of ShapeNet cars (see "Generating the ShapeNet datasets" for additional details).

For the following source object and view:

and target view:

run_pixelnerf.py generated the following after 73,243 iterations (~12 hours on a P100 GPU; the full pixelNeRF model was trained for 400,000 iterations, which took six days):

Loss: 0.004468636587262154

The "smearing" is an artifact caused by the bounding box sampling method.

Generating the ShapeNet datasets

1. Download the data (the ShapeNet server is pretty slow, so this will take a while):

SHAPENET_BASE_DIR=<path/to/your/shapenet/root>
nohup wget --quiet -P ${SHAPENET_BASE_DIR} http://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v2.zip > shapenet.log &

2. Unzip the data:

cd ${SHAPENET_BASE_DIR}
nohup unzip -q ShapeNetCore.v2.zip > shapenet.log &

3. After the file is done unzipping, remove the ZIP:

rm ShapeNetCore.v2.zip

4. Change the SHAPENET_DIR variable in generate_nerf_dataset.py and generate_pixelnerf_dataset.py to <path/to/your/shapenet/root/ShapeNetCore.v2>.