daniel03c1 / masked_wavelet_nerf

Masked Wavelet Representation for Compact Neural Radiance Fields

Daniel Rho*, Byeonghyeon Lee*, Seungtae Nam, Joo Chan Lee, Jong Hwan Ko†, and Eunbyung Park†, CVPR 2023

Project Page, Paper(arxiv)

Our code is based on TensoRF (https://github.com/apchenstu/TensoRF).

Our method, however, can be applied to any 2D grid-based neural fields.

Tested on Ubuntu 18.04 + Pytorch 1.10.2

0. Requirements

conda create -n MaskDWT python=3.8
conda activate MaskDWT
pip install torch torchvision
pip install tqdm scikit-image opencv-python configargparse lpips imageio-ffmpeg kornia lpips tensorboard

0.1 Installing pytorch_wavelets

git clone https://github.com/fbcotter/pytorch_wavelets
cd pytorch_wavelets
pip install .

1. Training

python3 train.py --config=configs/chair.txt --use_mask --mask_weight=1e-10 --grid_bit=8 --use_dwt --dwt_level=4

• "--config": the name of the config file
• "--datadir": the directory of images
• "--grid_bit": for n-bit quantization (QAT) (only works on grid parameters)
• "--use_mask": in order to use trainable masks, use this option
• "--mask_weight": loss weight to modulate the grid sparsity
• "--use_dwt": to use the wavelet transform
• "--dwt_level": the level of wavelet transform
• "--trans_func": the name of the wavelet function (default=bior4.4)

More details can be found in "opt.py"

2. Model Compression

python3 compress.py --compress=1 --compress_levelwise=1 --ckpt=PATH_TO_CHECKPOINT

• "--ckpt": the saved file name
• "--compress": set to a non-zero value for compression (default=0)
• "--compress_levelwise": set to a non-zero value for level-wise compression (default=0)
• "--decompress": set to a non-zero value for compression (default=0)
• "--decompress_levelwise": you need to use the same value as "compress_levelwise." (default=0)
• "--decompress_and_validatd": whether to evaluate the quality of a decompressed model (default=1)

3. Decompression and Evaluation

python3 compress.py --decompress=1 --decompress_levelwise=1 --config=configs/chair.txt --ckpt=PATH_TO_CHECKPOINT

4. Bibtex

@InProceedings{Rho_2023_CVPR,
    author    = {Rho, Daniel and Lee, Byeonghyeon and Nam, Seungtae and Lee, Joo Chan and Ko, Jong Hwan and Park, Eunbyung},
    title     = {Masked Wavelet Representation for Compact Neural Radiance Fields},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {20680-20690}
}