ZipNeRF

一个非官方的 PyTorch 实现，用于“抗锯齿基于网格的神经辐射场”（Anti-Aliased Grid-Based Neural Radiance Fields）https://arxiv.org/abs/2304.06706。 本项目基于 multinerf，因此 refnerf、rawnerf、mipnerf360 中的功能也均适用。

更新日志

• (2024.2.2) 添加对 nerfstudio 的支持，感谢 Ling Jing。
• (2024.12.8) 添加对 Intel 的 DPC++ 后端的支持，感谢 Zong Wei。
• (2023.6.22) 添加通过 tsdf 提取网格；为近景漂浮伪影添加 梯度缩放。
• (2023.5.26) 实现最新版本的 ZipNeRF https://arxiv.org/abs/2304.06706。
• (2023.5.22) 添加提取网格；添加日志记录、检查点系统

结果

新结果 (5.27): 预训练权重

360_v2:

https://github.com/SuLvXiangXin/zipnerf-pytorch/assets/83005605/2b276e48-2dc4-4508-8441-e90ec963f7d9

360_v2_glo:(更少的漂浮伪影，但指标较差)

https://github.com/SuLvXiangXin/zipnerf-pytorch/assets/83005605/bddb5610-2a4f-4981-8e17-71326a24d291

网格结果 (5.27):

MipNeRF360(峰值信噪比 (PSNR)):

	bicycle	garden	stump	room	counter	kitchen	bonsai
论文	25.80	28.20	27.55	32.65	29.38	32.50	34.46
本仓库	25.44	27.98	26.75	32.13	29.10	32.63	34.20

Blender(峰值信噪比 (PSNR)):

	chair	drums	ficus	hotdog	lego	materials	mic	ship
论文	34.84	25.84	33.90	37.14	34.84	31.66	35.15	31.38
本仓库	35.26	25.51	32.66	36.56	35.04	29.43	34.93	31.38

对于 MipNeRF360 数据集，模型使用 4 的下采样因子 (downsample factor) 训练室外场景，室内场景使用 2（与论文相同）。 训练速度约为论文的 1.5 倍慢（在 8 张 A6000 上耗时约 1.5 小时）。

Hash 衰减损失 (hash decay loss) 似乎影响很小 (?), 因为在两个实验的最终结果中都能发现许多漂浮伪影 (floaters)（尤其是在 Blender 中）。

安装 CUDA 后端

# Clone the repo.
git clone https://github.com/SuLvXiangXin/zipnerf-pytorch.git
cd zipnerf-pytorch

# Make a conda environment.
conda create --name zipnerf python=3.9
conda activate zipnerf

# Install requirements.
pip install -r requirements.txt

# Install other cuda extensions
pip install ./extensions/cuda

# Install nvdiffrast (optional, for textured mesh)
git clone https://github.com/NVlabs/nvdiffrast
pip install ./nvdiffrast

# Install a specific cuda version of torch_scatter 
# see more detail at https://github.com/rusty1s/pytorch_scatter
CUDA=cu117
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.html

安装 DPCPP 后端

  # Install drivers, oneAPI and ipex for Intel GPUs
  Following the steps in the below page to install gpu drivers, oneAPI BaseKit, and pytorch+ipex (abbr. intel-extension-for-pytorch):
  https://intel.github.io/intel-extension-for-pytorch/xpu/1.13.120+xpu/tutorials/installation.html
  For pytorch and Ipex versions, please install the version 1.13.120 with

  python -m pip install torch==1.13.0a0+git6c9b55e intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu

  After the installation is done, make sure it is successfully by running the example provided by
  https://github.com/intel/intel-extension-for-pytorch/tree/release/xpu/1.13.120#inference-on-gpu

准备环境

  export DPCPP_HOME=path/to/llvm  # path to the folder for llvm, default value:~
  bash scripts/set_dpcpp_env.sh intel # for intel's gpu
  bash scripts/set_dpcpp_env.sh nvidia # for nvidia's gpu

DPCPP 支持 CUDA 参考

  https://github.com/intel/llvm/blob/sycl/sycl/doc/GetStartedGuide.md#build-dpc-toolchain-with-support-for-nvidia-cuda

数据集

mipnerf360

refnerf

nerf_synthetic

nerf_llff_data

mkdir data
cd data

# e.g. mipnerf360 data
wget http://storage.googleapis.com/gresearch/refraw360/360_v2.zip
unzip 360_v2.zip

训练

# Configure your training (DDP? fp16? ...)
# see https://huggingface.co/docs/accelerate/index for details
accelerate config

# Where your data is 
DATA_DIR=data/360_v2/bicycle
EXP_NAME=360_v2/bicycle

# Experiment will be conducted under "exp/${EXP_NAME}" folder
# "--gin_configs=configs/360.gin" can be seen as a default config 
# and you can add specific config useing --gin_bindings="..." 
accelerate launch train.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.factor = 4"

# or you can also run without accelerate (without DDP)
CUDA_VISIBLE_DEVICES=0 python train.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
      --gin_bindings="Config.factor = 4" 

# alternatively you can use an example training script 
bash scripts/train_360.sh

# blender dataset
bash scripts/train_blender.sh

# metric, render image, etc can be viewed through tensorboard
tensorboard --logdir "exp/${EXP_NAME}"

使用 DPCPP 后端进行训练和渲染

# add config in command line
      --gin_bindings="Config.dpcpp_backend = True" \

渲染

渲染结果可在目录 exp/${EXP_NAME}/render 中找到

accelerate launch render.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.render_path = True" \
    --gin_bindings="Config.render_path_frames = 480" \
    --gin_bindings="Config.render_video_fps = 60" \
    --gin_bindings="Config.factor = 4"  

# alternatively you can use an example rendering script 
bash scripts/render_360.sh

评估

评估结果可在目录 exp/${EXP_NAME}/test_preds 中找到

# using the same exp_name as in training
accelerate launch eval.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.factor = 4"


# alternatively you can use an example evaluating script 
bash scripts/eval_360.sh

使用 NerfStudio

https://github.com/nerfstudio-project/nerfstudio
NerfStudio 提供了一个简单的 API（应用程序编程接口），允许简化创建、训练和测试 NeRF（神经辐射场）的端到端流程。该库通过模块化每个组件来支持更可解释的 NeRF 实现。你可以使用 NerfStudio 提供的查看器在训练过程中查看渲染结果。

安装

pip install nerfstudio  
# cd zipnerf-pytorch
pip install -e . 
ns-install-cli

训练与评估

ns-train zipnerf --data {DATA_DIR/SCENE}
ns-eval --load-config {outputs/SCENE/zipnerf/EXP_DIR/config.yml}

ns-train zipnerf -h  # show the full list of model configuration options.
ns-train zipnerf colmap -h  # dataparset configuration options

*NerfStudio 的 ColmapDataParser 在下采样时会向下取整图像尺寸，这与 360_v2 数据集不同。你可以使用 NerfStudio 重新处理数据，或者按照 https://github.com/nerfstudio-project/nerfstudio/issues/1438 中讨论的那样修改库中的下采样代码逻辑。
*NerfStudio 的训练/评估划分策略与此仓库不同。最终的训练和评估结果可能会有所差异。

更多用法或信息，请参阅 https://github.com/nerfstudio-project/nerfstudio。

配置

针对 Zipnerf-pytorch

你可以创建一个新的 .gin 文件，并在 zipnerf_ns/zipnerf_config.py 的 ZipNerfModelConfig 中传入 'gin_file' 列表，或者更新默认 .gin 文件的内容。

针对 NerfStudio

ns-train zipnerf -h
ns-train zipnerf colmap -h

你可以修改 zipnerf_ns/zipnerf_config.py，或者使用上述指令。

查看器

给定一个预训练的模型检查点，你可以通过运行以下命令启动查看器：

ns-viewer --load-config outputs/SCENE/zipnerf/EXP_TIME/config.yml  

远程服务器

如果你在远程机器上运行，你需要进行 WebSocket 端口转发（默认为 7007）。必须在远程机器上设置 SSH。然后在此机器上运行以下内容：

ssh -L <port>:localhost:<port> USER@REMOTE.SERVER.IP

提取网格

网格结果可以在目录 exp/${EXP_NAME}/mesh 中找到

# more configuration can be found in internal/configs.py
accelerate launch extract.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.factor = 4"
#    --gin_bindings="Config.mesh_radius = 1"  # (optional) smaller for more details e.g. 0.2 in bicycle scene
#    --gin_bindings="Config.isosurface_threshold = 20"  # (optional) empirical value
#    --gin_bindings="Config.mesh_voxels=134217728"  # (optional) number of voxels used to extract mesh, e.g. 134217728 equals to 512**3 . Smaller values may solve OutoFMemoryError
#    --gin_bindings="Config.vertex_color = True"  # (optional) saving mesh with vertex color instead of atlas which is much slower but with more details.
#    --gin_bindings="Config.vertex_projection = True"  # (optional) use projection for vertex color

# or extracting mesh using tsdf method
accelerate launch tsdf.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.factor = 4"

# alternatively you can use an example script 
bash scripts/extract_360.sh

内存不足

你可以通过添加例如 --gin_bindings="Config.batch_size = 8192" 来减少总批次大小，或者通过添加例如 --gin_bindings="Config.render_chunk_size = 8192" 来减少测试分块大小，或者通过配置 accelerate config 使用更多的 GPU。

准备自定义数据

更多详细信息可在 https://github.com/google-research/multinerf 找到

DATA_DIR=my_dataset_dir
bash scripts/local_colmap_and_resize.sh ${DATA_DIR}

待办事项

• 添加多尺度训练和测试

引用

@misc{barron2023zipnerf,
      title={Zip-NeRF: Anti-Aliased Grid-Based Neural Radiance Fields}, 
      author={Jonathan T. Barron and Ben Mildenhall and Dor Verbin and Pratul P. Srinivasan and Peter Hedman},
      year={2023},
      eprint={2304.06706},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

@misc{multinerf2022,
      title={{MultiNeRF}: {A} {Code} {Release} for {Mip-NeRF} 360, {Ref-NeRF}, and {RawNeRF}},
      author={Ben Mildenhall and Dor Verbin and Pratul P. Srinivasan and Peter Hedman and Ricardo Martin-Brualla and Jonathan T. Barron},
      year={2022},
      url={https://github.com/google-research/multinerf},
}

@Misc{accelerate,
  title =        {Accelerate: Training and inference at scale made simple, efficient and adaptable.},
  author =       {Sylvain Gugger, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Mangrulkar},
  howpublished = {\url{https://github.com/huggingface/accelerate}},
  year =         {2022}
}

@misc{torch-ngp,
    Author = {Jiaxiang Tang},
    Year = {2022},
    Note = {https://github.com/ashawkey/torch-ngp},
    Title = {Torch-ngp: a PyTorch implementation of instant-ngp}
}

致谢

这项工作基于我的另一个仓库 https://github.com/SuLvXiangXin/multinerf-pytorch，它基本上是从 multinerf 进行的 PyTorch 移植。

• 感谢 multinerf 提供了出色的 multinerf(MipNeRF360,RefNeRF,RawNeRF) 实现
• 感谢 accelerate 提供的分布式训练支持
• 感谢 torch-ngp 提供的超有用的 hashencoder
• 感谢 Yurui Chen 对论文细节的讨论。

zipnerf-pytorch 快速上手指南

ZipNeRF 是论文 "Zip-NeRF: Anti-Aliased Grid-Based Neural Radiance Fields" 的非官方 PyTorch 实现。该工具基于 multinerf，支持训练、渲染、网格提取及 Nerfstudio 集成，适用于神经辐射场（NeRF）相关研究与应用。

环境准备

• 操作系统: Linux / Windows (推荐 Linux)
• Python: 3.9+
• GPU: NVIDIA GPU (CUDA 后端) 或 Intel GPU (DPC++ 后端)
• 依赖管理: 推荐使用 Conda 管理虚拟环境
• 前置软件: Git, CUDA Toolkit (若使用 CUDA 后端)

安装步骤

1. 克隆仓库与创建环境

git clone https://github.com/SuLvXiangXin/zipnerf-pytorch.git
cd zipnerf-pytorch

conda create --name zipnerf python=3.9
conda activate zipnerf

2. 安装基础依赖

pip install -r requirements.txt

3. 安装 CUDA 扩展 (核心步骤)

pip install ./extensions/cuda

4. 安装可选组件

• 纹理网格提取 (需要 nvdiffrast):

  git clone https://github.com/NVlabs/nvdiffrast
  pip install ./nvdiffrast
  

• torch_scatter (需匹配当前 CUDA 版本，例如 cu117):

  CUDA=cu117
  pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.html
  

注意: 如需使用 Intel GPU (DPC++ 后端)，请参考 README 中的 Install DPCPP backend 章节配置 OneAPI 和 ipex。

基本使用

1. 数据集准备

下载官方数据集并解压至 data 目录：

mkdir data
cd data
# 以 mipnerf360 为例
wget http://storage.googleapis.com/gresearch/refraw360/360_v2.zip
unzip 360_v2.zip

2. 配置加速环境

在使用 accelerate 进行分布式训练前，需先配置：

accelerate config

3. 模型训练

以 bicycle 场景为例，启动训练：

DATA_DIR=data/360_v2/bicycle
EXP_NAME=360_v2/bicycle

accelerate launch train.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.factor = 4"

查看训练进度可使用 TensorBoard：

tensorboard --logdir "exp/${EXP_NAME}"

4. 渲染视频

训练完成后，生成漫游视频：

accelerate launch render.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.render_path = True" \
    --gin_bindings="Config.render_path_frames = 480" \
    --gin_bindings="Config.render_video_fps = 60" \
    --gin_bindings="Config.factor = 4"

结果保存在 exp/${EXP_NAME}/render 目录。

5. 网格提取 (Mesh)

提取三维网格模型：

accelerate launch extract.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = '${DATA_DIR}'" \
    --gin_bindings="Config.exp_name = '${EXP_NAME}'" \
    --gin_bindings="Config.factor = 4"

结果保存在 exp/${EXP_NAME}/mesh 目录。

6. Nerfstudio 集成 (可选)

支持通过 Nerfstudio 进行更便捷的端到端流程：

pip install nerfstudio
pip install -e .
ns-install-cli

# 训练
ns-train zipnerf --data {DATA_DIR/SCENE}

# 评估
ns-eval --load-config {outputs/SCENE/zipnerf/EXP_DIR/config.yml}

# 可视化
ns-viewer --load-config outputs/SCENE/zipnerf/EXP_TIME/config.yml

---

提示: 如遇显存不足 (OutOfMemory)，可通过减小 batch_size 或 render_chunk_size 解决，例如添加 --gin_bindings="Config.batch_size = 8192"。