ZipNeRF (modernized)

A PyTorch implementation of Zip-NeRF: Anti-Aliased Grid-Based Neural Radiance Fields (Barron et al., 2023).

This repository is a modernized fork of SuLvXiangXin/zipnerf-pytorch. The upstream code was written against Python 3.9, PyTorch 2.0, CUDA 11.7, and NumPy 1.x and no longer builds on current toolchains. This fork tracks the same algorithm but is patched to compile and run on the 2026 stack.

What was changed vs. upstream

Component	Upstream	This fork
Python	3.9	3.11+
PyTorch	2.0	2.9+
CUDA	11.7	12.8
NumPy	1.x	2.x
GPU arch	sm_70+	sm_70+ incl. Blackwell sm_120 (RTX 50-series)

Specific source changes:

CUDA extension (extensions/cuda/): bumped to -std=c++17 (required by PyTorch 2.x); fixed the atomicAdd(at::Half*, at::Half) stub that was missing a return statement and silently produced garbage on newer nvcc.
Python AMP API (gridencoder/grid.py, extensions/backend.py): migrated torch.cuda.amp.{custom_fwd,custom_bwd,autocast} to the new torch.amp API with device_type='cuda', required since PyTorch 2.4.
pycolmap (internal/pycolmap/pycolmap/scene_manager.py): np.uint64(-1) → np.iinfo(np.uint64).max. NumPy 2.x no longer silently wraps negative ints into unsigned types.
nerfstudio integration removed to keep the dependency surface small. If you need it, check out upstream.

Install

Tested on Ubuntu 22 / WSL2, CUDA 12.8, PyTorch 2.9, Python 3.11.

# 1. Clone the repo
git clone https://github.com/HyoKong/zipnerf.git
cd zipnerf

# 2. Create a conda environment (any name; we use "zipnerf" here)
conda create -n zipnerf python=3.11 -y
conda activate zipnerf

# 3. Install PyTorch first, matching your CUDA. Example for CUDA 12.8:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128

# 4. Other Python dependencies
pip install -r requirements.txt

# 5. Build the custom CUDA backend (hash grid + PDF sampler)
pip install --no-build-isolation ./extensions/cuda

# 6. torch_scatter — IMPORTANT: build from source so it matches your torch/CUDA.
#    The cached wheels on PyPI are often ABI-incompatible with new torch builds.
pip install --no-build-isolation --no-binary=torch-scatter --no-cache-dir torch-scatter

Quick sanity check:

python -c "
from extensions import Backend; Backend.set_backend('cuda')
from gridencoder import GridEncoder
import torch
enc = GridEncoder(input_dim=3, num_levels=4, level_dim=2, base_resolution=16).cuda()
x = (torch.rand(1024, 3, device='cuda') * 2 - 1).requires_grad_(True)
y = enc(x); y.sum().backward()
print('OK:', y.shape, x.grad.shape)
"

Dataset

The repo's defaults assume mip-NeRF 360 v2 is unzipped under data/360_v2/. The dataset ships as two zips — the main 7 scenes plus 2 extras (flowers, treehill) needed for the full 9-scene benchmark used by the paper.

mkdir -p data/360_v2 && cd data/360_v2

# 7 main scenes (~9 GB)
wget http://storage.googleapis.com/gresearch/refraw360/360_v2.zip
unzip 360_v2.zip && rm 360_v2.zip

# 2 extra scenes (~6 GB) — flowers + treehill
wget http://storage.googleapis.com/gresearch/refraw360/360_extra_scenes.zip
unzip 360_extra_scenes.zip && rm 360_extra_scenes.zip

cd ../..

Resulting layout (per scene):

data/360_v2/<scene>/
  images/        # original resolution
  images_2/      # downsampled ×2
  images_4/      # downsampled ×4
  images_8/      # downsampled ×8
  poses_bounds.npy
  sparse/        # COLMAP output

The 9 standard scenes:

Outdoor (factor=4)	Indoor (factor=2)
`bicycle`, `garden`, `stump`, `flowers`, `treehill`	`room`, `counter`, `kitchen`, `bonsai`

If you'd rather use a different path, every command in this README and the multi-GPU script accepts a custom location (override Config.data_dir for single scenes, or DATA_ROOT for the batch script).

Other supported datasets (see configs/):

Train a single scene

Outdoor uses factor=4, indoor uses factor=2 (matches the paper). The default batch_size=65536 assumes high-memory GPUs (≥24 GB). Drop it for smaller cards (see Memory tuning).

5k-step trial (~13 min on RTX 5070 Ti, useful to validate the setup):

CUDA_VISIBLE_DEVICES=0 python train.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = 'data/360_v2/bicycle'" \
    --gin_bindings="Config.exp_name = '360_v2/bicycle'" \
    --gin_bindings="Config.factor = 4" \
    --gin_bindings="Config.batch_size = 8192" \
    --gin_bindings="Config.render_chunk_size = 8192" \
    --gin_bindings="Config.max_steps = 5000"

Full 25k-step run (paper recipe):

CUDA_VISIBLE_DEVICES=0 python train.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = 'data/360_v2/bicycle'" \
    --gin_bindings="Config.exp_name = '360_v2/bicycle'" \
    --gin_bindings="Config.factor = 4"

Outputs:

exp/360_v2/bicycle/checkpoints/<step>/ — model + optimizer state
exp/360_v2/bicycle/events.out.tfevents.* — TensorBoard log
exp/360_v2/bicycle/log_train.txt — training log

Live monitoring:

tensorboard --logdir exp/360_v2

Evaluate

Computes PSNR / SSIM / LPIPS over the held-out test split (every 8th image when llffhold=8, the default for LLFF/360):

CUDA_VISIBLE_DEVICES=0 python eval.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = 'data/360_v2/bicycle'" \
    --gin_bindings="Config.exp_name = '360_v2/bicycle'" \
    --gin_bindings="Config.factor = 4"

Per-scene output, all under exp/360_v2/bicycle/test_preds/:

File	Contents
`metrics.csv`	per-image rows + final mean: PSNR / SSIM / LPIPS (raw and color-corrected). Stable name pointing to the latest checkpoint.
`metrics_<step>.csv`	same data, stamped with the checkpoint step.
`color_<idx>.png`	rendered RGB.
`color_cc_<idx>.png`	rendered RGB color-corrected to GT (used for `*_cc` metrics).
`gt_<idx>.png`	matched ground-truth, saved for direct comparison.
`compare_<idx>.png`	side-by-side `[GT \| render]` (4 px black separator) for one-glance comparison in a file browser.
`distance_mean_<idx>.tiff`, `distance_median_<idx>.tiff`	raw depth (float32).
`distance_mean_viz_<idx>.png`	turbo-colormap depth for quick browsing.
`acc_<idx>.tiff`	accumulated alpha.
`metric_*_<step>.txt`	legacy per-metric space-separated dumps.

Render a turntable video

CUDA_VISIBLE_DEVICES=0 python render.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = 'data/360_v2/bicycle'" \
    --gin_bindings="Config.exp_name = '360_v2/bicycle'" \
    --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"

Output in exp/360_v2/bicycle/render/.

Train + eval every 360 scene with one command

scripts/train_all_360.sh loops over all nine standard mip-NeRF 360 scenes (sequential), and for each scene runs train → eval using all visible GPUs in DDP. The script keeps going if one scene fails (per-scene failure does not abort the rest).

# defaults: auto-detect GPUs, 25k steps, batch=65536 global, train + eval each scene
bash scripts/train_all_360.sh

# explicit 4-GPU run, full benchmark
NUM_GPUS=4 bash scripts/train_all_360.sh

# subset of scenes, quick trial
SCENES="bicycle garden" MAX_STEPS=5000 BATCH_SIZE=32768 \
    bash scripts/train_all_360.sh

# skip eval (train only)
SKIP_EVAL=1 bash scripts/train_all_360.sh

Configurable env vars:

Var	Default	Meaning
`NUM_GPUS`	auto	GPUs per scene (DDP via accelerate)
`SCENES`	all 9 standard 360 scenes	space-separated subset
`MAX_STEPS`	25000	training iterations per scene
`BATCH_SIZE`	65536	global batch (auto-split across GPUs)
`RENDER_CHUNK_SIZE`	= `BATCH_SIZE`	eval chunk size
`DATA_ROOT`	`data/360_v2`	where the scene folders live
`EXP_PREFIX`	`360_v2`	output goes to `exp/${EXP_PREFIX}/<scene>/`
`SKIP_EVAL`	`0`	set to `1` to train only and skip eval

The script auto-selects factor=4 for outdoor scenes (bicycle / garden / stump / flowers / treehill) and factor=2 for indoor (room / counter / kitchen / bonsai). It gracefully skips scenes whose folder is missing — useful if you've only downloaded part of the dataset.

Per scene you'll get a populated exp/${EXP_PREFIX}/<scene>/test_preds/ directory containing metrics.csv and the rendered / ground-truth / depth images described in the Evaluate section above.

Extract a mesh

# marching-cubes from the density field
CUDA_VISIBLE_DEVICES=0 python extract.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = 'data/360_v2/bicycle'" \
    --gin_bindings="Config.exp_name = '360_v2/bicycle'" \
    --gin_bindings="Config.factor = 4"

# or TSDF fusion from rendered depths
CUDA_VISIBLE_DEVICES=0 python tsdf.py \
    --gin_configs=configs/360.gin \
    --gin_bindings="Config.data_dir = 'data/360_v2/bicycle'" \
    --gin_bindings="Config.exp_name = '360_v2/bicycle'" \
    --gin_bindings="Config.factor = 4"

Memory tuning

batch_size is the number of rays per training step; render_chunk_size is the number of rays processed at once during eval/render (memory tiling, mathematically identical to a full pass).

GPU VRAM	Suggested `batch_size`	Notes
16 GB (5070 Ti, 4080)	8192 – 16384	verified on 5070 Ti
24 GB (3090, 4090)	16384 – 32768
48 GB (A6000, A100 40 GB)	32768 – 65536
80 GB (A100 80 GB, H100)	65536	paper default

Set render_chunk_size to the same scale as batch_size.

System RAM rises during the first ~10 s of training because the dataloader caches every training image. For a standard 360 scene at factor=4, expect ~2–3 GB resident; at factor=2, ~8–10 GB.

Custom data

The upstream COLMAP helper still works:

DATA_DIR=path/to/my_dataset
bash scripts/local_colmap_and_resize.sh ${DATA_DIR}

Then train it like any other scene (Config.dataset_loader = 'llff').

Reproduced results

Two configurations:

8× V100 — paper recipe: max_steps=25000, batch_size=65536 (global, split 8 ways), all other hyper-parameters at upstream defaults. This is the apples-to-apples reproduction.
1× RTX 5070 Ti — same recipe except batch_size=8192 (the largest that fits in 16 GB). Documents what you get on a consumer GPU.

Numbers are raw PSNR / SSIM on the held-out 1/8 LLFF split (color-corrected metrics are in each scene's test_preds/metrics.csv).

mip-NeRF 360, 25k steps — PSNR

	bicycle	garden	stump	flowers	treehill	room	counter	kitchen	bonsai	mean (7)
Paper (7)	25.80	28.20	27.55	—	—	32.65	29.38	32.50	34.46	30.08
Upstream (7)	25.44	27.98	26.75	—	—	32.13	29.10	32.63	34.20	29.75
This fork — 8× V100	25.38	27.99	26.70	22.12	23.17	32.03	28.98	32.57	34.26	29.70
This fork — 1× 5070 Ti (b=8192)	24.57	26.84	26.30	21.57	23.25	31.38	28.50	31.82	32.91	28.90

mip-NeRF 360, 25k steps — SSIM

	bicycle	garden	stump	flowers	treehill	room	counter	kitchen	bonsai
This fork — 8× V100	0.772	0.868	0.784	0.646	0.645	0.927	0.903	0.937	0.954
This fork — 1× 5070 Ti (b=8192)	0.697	0.824	0.752	0.591	0.605	0.913	0.886	0.926	0.943

The paper and the upstream repo only publish numbers for the main 7 scenes, so flowers and treehill have no comparison row. The mean (7) column averages over the 7 scenes that all four rows have, for an apples-to-apples comparison.

Notes:

The 8× V100 row matches upstream within 0.05 dB mean PSNR (per-scene deltas all within 0.12 dB; bonsai and garden slightly beat upstream). At identical hyper-parameters and 25 k steps, this is a faithful reproduction — it confirms the modernization patches (Torch 2.9 / CUDA 12.8 / NumPy 2.x) don't perturb training numerics.
The 1× 5070 Ti row trails the V100 row by ~0.8 dB mean PSNR. The single change is batch_size=8192 vs 65 536; this is the expected cost of an 8× batch-size reduction (smaller effective sample count over the same 25 k steps and the LR schedule wasn't retuned). See Memory tuning for the VRAM ↔ batch table.
If you can run multi-GPU, just NUM_GPUS=8 bash scripts/train_all_360.sh reproduces the V100 row.
Per-image CSVs, paired GT + rendered images (compare_<idx>.png), and colorized depth visualizations are saved under exp/360_v2/<scene>/test_preds/ (see the Evaluate section for the file layout).

Troubleshooting

undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKSsb when importing torch_scatter: the cached PyPI wheel is ABI-incompatible. Rebuild from source:
```
pip uninstall -y torch-scatter
pip install --no-build-isolation --no-binary=torch-scatter --no-cache-dir torch-scatter
```
ModuleNotFoundError: No module named 'torch' during pip install ./extensions/cuda: PEP 517 build isolation hides your env's torch. Use --no-build-isolation (covered in install step 5).
CUDA OOM on first training step: drop Config.batch_size (and Config.render_chunk_size) — see the memory table above.
OverflowError: Python integer -1 out of bounds for uint64 when loading COLMAP data: fixed in this fork; pull the latest.

Citation

The algorithm is from:

@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}
}

Original codebase this fork extends:

SuLvXiangXin/zipnerf-pytorch — unofficial PyTorch port
google-research/multinerf — reference MipNeRF360 / RefNeRF / RawNeRF
Yurui Chen and Jiaxiang Tang's torch-ngp — hash encoder

License

See LICENSE. All rights to the original work remain with the original authors.

Navigation Menu

Repository files navigation

ZipNeRF (modernized)

What was changed vs. upstream

Install

Dataset

Train a single scene

Evaluate

Render a turntable video

Train + eval every 360 scene with one command

Extract a mesh

Memory tuning

Custom data

Reproduced results

mip-NeRF 360, 25k steps — PSNR

mip-NeRF 360, 25k steps — SSIM

Troubleshooting

Citation

License

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
Latest commit History 77 Commits
configs		configs
extensions		extensions
gridencoder		gridencoder
internal		internal
scripts		scripts
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
eval.py		eval.py
extract.py		extract.py
render.py		render.py
requirements.txt		requirements.txt
train.py		train.py
tsdf.py		tsdf.py

Folders and files

Latest commit

History

Repository files navigation

ZipNeRF (modernized)

What was changed vs. upstream

Install

Dataset

Train a single scene

Evaluate

Render a turntable video

Train + eval every 360 scene with one command

Extract a mesh

Memory tuning

Custom data

Reproduced results

mip-NeRF 360, 25k steps — PSNR

mip-NeRF 360, 25k steps — SSIM

Troubleshooting

Citation

License

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages