Skip to content

alexanderbergman7/metanlrpp

Repository files navigation

Fast Training of Neural Lumigraph Representations using Meta Learning

Alexander W. Bergman, Petr Kellnhofer, Gordon Wetzstein, Stanford University.
Official Implementation for Fast Training of Neural Lumigraph Representations using Meta Learning.

Usage

To get started, create a conda environment with all dependencies:

conda env create -f environment.yml
conda activate metanlrpp

Code Structure

The code is organized as follows:

  • experiment_scripts: directory containing scripts to for training and testing MetaNLR++ models.
    • pretrain_features.py: pre-train encoder and decoder networks
    • train_sdf_ibr_meta.py: train meta-learned initialization for encoder, decoder, aggregation fn, and neural SDF
    • test_sdf_ibr_meta.py: specialize meta-learned initialization to a specific scene
    • train_sdf_ibr.py: train NLR++ model from scratch without meta-learned initialization
    • test_sdf_ibr.py: evaluate performance on withheld views
  • configs: directory containing configs to reproduce experiments in the paper
    • nlrpp_nlr.txt: configuration for training NLR++ on the NLR dataset
    • nlrpp_dtu.txt: configuration for training NLR++ on the DTU dataset
    • nlrpp_nlr_meta.txt: configuration for training the MetaNLR++ initialization on the NLR dataset
    • nlrpp_dtu_meta.txt: configuration for training the MetaNLR++ initialization on the DTU dataset
    • nlrpp_nlr_metaspec.txt: configuration for training MetaNLR++ on the NLR dataset using the learned initialization
    • nlrpp_dtu_metaspec.txt: configuration for training MetaNLR++ on the DTU dataset using the learned initialization
  • data_processing: directory containing utility functions for processing data
  • torchmeta: torchmeta library for meta-learning
  • utils: directory containing various utility functions for rendering and visualization
  • loss_functions.py: file containing loss functions for evaluation
  • meta_modules.py: contains meta learning wrappers around standard modules using torchmeta
  • modules.py: contains standard modules for coodinate-based networks
  • modules_sdf.py: extends standard modules for coordinate-based network representations of signed-distance functions.
  • modules_unet.py: contains encoder and decoder modules used for image-space feature processing
  • scheduler.py: utilities for training schedule
  • training.py: training script
  • sdf_rendering.py: functions for rendering SDF
  • sdf_meshing.py: functions for meshing SDF
  • checkpoints: contains checkpoints to some pre-trained models (additional/ablation models by request)
  • assets: contains paths to checkpoints which are used as assets, and pre-computed buffers over multiple runs (if necessary)

Getting Started

Pre-training Encoder and Decoder

Pre-train the encoder and decoder using the FlyingChairsV2 training dataset as follows:

python experiment_scripts/pretrain_features.py --experiment_name XXX --batch_size X --dataset_path /path/to/FlyingChairs2/train

Alternatively, use the checkpoint in the checkpoints directory.

Training NLR++

Train a NLR++ model using the following command:

python experiment_scripts/train_sdf_ibr.py --config_filepath configs/nlrpp_dtu.txt --experiment_name XXX --dataset_path /path/to/dtu/scanXXX --checkpoint_img_encoder /path/to/pretrained/encdec

Note that we have uploaded our processed version of the DTU and NLR data here. The raw NLR data can be found here.

Meta-learned Initialization (MetaNLR++)

Meta-learn the initialization for the encoder, decoder, aggregation function, and neural SDF using the following command:

python experiment_scripts/train_sdf_ibr_meta.py --config_filepath configs/nlrpp_dtu_meta.txt --experiment_name XXX --dataset_path /path/to/dtu/meta/training --reference_view 24 --checkpoint_img_encoder /path/to/pretrained/encdec

Some optimized initializations for the DTU and NLR datasets can be found in the data directory. Additional models can be provided upon request.

Training MetaNLR++ from Initialization

Use the meta-learned initialization to specialize to a specific scene using the following command:

python experiment_scripts/test_sdf_ibr_meta.py --config_filepath configs/nlrpp_dtu_metaspec.txt --experiment_name XXX --dataset_path /path/to/dtu/scanXXX --reference_view 24 --meta_initialization /path/to/learned/meta/initialization

Evaluation

Test the converged scene on withheld views using the following command:

python experiment_scripts/test_sdf_ibr.py --config_filepath configs/nlrpp_dtu.txt --experiment_name XXX --dataset_path /path/to/dtu/scanXXX --checkpoint_path_test /path/to/checkpoint/to/evaluate

Citation & Contact

If you find our work useful in your research, please cite

@inproceedings{bergman2021metanlr,
author = {Bergman, Alexander W. and Kellnhofer, Petr and Wetzstein, Gordon},
title = {Fast Training of Neural Lumigraph Representations using Meta Learning},
booktitle = {NeurIPS},
year = {2021},
}

If you have any questions or would like access to specific ablations or baselines presented in the paper or supplement (the code presented here is only a subset based off of the source code used to generate the results), please feel free to contact the authors. Alex can be contacted via e-mail at awb@stanford.edu.

About

Official Implementation for Fast Training of Neural Lumigraph Representations using Meta Learning.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages