dogyoonlee/nerf-pytorch-multi-gpu

Refactored PyTorch code of NeRF using multi gpu.

Pytorch Refactoring of NeRF for Multi-GPU.

Refactored nerf code based on pytorch to use the multi gpu.

We simply use nn.DataParallel based on the [DeblurNeRF] and [DP-NeRF] code referring the [nerf-pytorch] code.

Training & Evaluation

1. Environment

git clone https://github.com/dogyoonlee/nerf-pytorch-multi-gpu.git
cd nerf-pytorch-multi-gpu
pip install -r requirements.txt

Dependencies (click to expand)

numpy

scikit-image

torch>=1.11

torchvision>=0.12.0

imageio

imageio-ffmpeg

matplotlib

configargparse

tensorboardX>=2.0

opencv-python

einops

tensorboard

python-box

pyyaml

tqdm

mediapy

pillow

Note that, you should match the version of torch and torchvision to properly work on your gpu.

Refer the torch install comments in requirements.txt.

2. Download dataset

There are various datasets based on NeRF

• NeRF dataset: [Github] [Dataset]

• HDR-NeRF dataset: [Github] [Dataset]

• Deblur-NeRF dataset: [Github] [Dataset]

We provide config example in configs folder.

Create your own config file referring the provided file.

3. Set parameters

Change the training parameters in config file to train on your computer.

You must set the config and expname parameter to run the training.

4. Train

For example, to train lego scene in nerf_synthetic,

python run_nerf.py --config ./configs/nerf_synthetic/tx_nerf_synthetic_lego.txt --expname <experiment_name>

The training and tensorboard results will be save in <basedir>/<expname> and <tbdir>.

5. Evaluation

Evaluation is automatically executed every --i_testset iterations. Please refer the other logging options in configs.py to adjust save and print the results.

After the training, execute the evaluation results following command. For example, to evaluate lego scene after 200000 iteration,

python run_nerf.py --config ./configs/nerf_synthetic/tx_nerf_synthetic_lego.txt --expname <dir_to_log> --ft_path ./<basedir>/<expname_trained>/200000.tar --render_only --render_test

Visualization of trained model

You can render or save the results after 200000 iteration of training following process.

1. Visualize the trained model as videos of spiral-path.

Results will be saved in ./<basedir>/<dir_to_log>/renderonly_path_199999.

python run_nerf.py --config ./configs/nerf_synthetic/tx_nerf_synthetic_lego.txt --expname <dir_to_log> --ft_path ./<basedir>/<expname_trained>/200000.tar --render_only 

2. Visualize the trained model as images of training views in datasets.

Results will be saved in ./<basedir>/<dir_to_log>/renderonly_test_199999

python run_dpnerf.py --config ./configs/nerf_synthetic/tx_nerf_synthetic_lego.txt --expname <dir_to_log> --ft_path ./<basedir>/<expname_trained>/200000.tar --render_only  --render_test

Notes

1. The number of GPU

If you want to set the number of GPU, change num_gpu value.

2. GPU memory

If you have not enough memory to train on your single GPU, set N_rand to a smaller value, or use multiple GPUs.

License

MIT License

References/Contributions

• nerf-official Code

• nerf-pytorch Code

• Deblur-NeRF Code

• DP-NeRF Code