使用JAX3D构建分层NeRF实现体渲染与三维重建教程
Hierarchical NeRF with JAX3D for Volumetric Rendering, Novel-View Synthesis, and 3D Reconstruction
推荐理由
这是一篇极具实操价值的硬核工程教程，给出了基于 JAX3D 实现分层 NeRF 的完整可复现工作流与关键参数配置，适合需要深入理解体渲染原理与落地 3D 重建的同学收藏参考。
In this tutorial, we build an end-to-end hierarchical Neural Radiance Field (NeRF) using JAX, Flax, Optax, and the volume-rendering primitives provided by jax3d. We first construct a synthetic multi-view dataset from an analytic scene containing volumetric geometry and view-dependent radiance, using sample_along_rays and volume_rendering to establish the forward rendering process. We then implement a NeRF with positional encoding, skip connections, separate coarse and fine networks, and view-direction conditioning, followed by hierarchical importance sampling through sample_piecewise_constant_pdf. We train the model with JAX JIT compilation, Adam optimization, exponential learning-rate decay, and gradient clipping, and finally evaluate novel-view synthesis using PSNR, depth and opacity visualization, sampling diagnostics, 360-degree rendering, and marching-cubes geometry extraction.
在本教程中，我们使用 JAX、Flax、Optax 以及 jax3d 提供的体积渲染原语，构建一个端到端的分层神经辐射场（NeRF）。首先，我们从包含体积几何结构和视角相关辐射的分析场景构建合成多视图数据集，利用 sample_along_rays 和 volume_rendering 建立前向渲染流程。接着，我们实现带有位置编码、跳跃连接、独立粗/细网络以及视角方向条件的 NeRF，并通过 sample_piecewise_constant_pdf 进行分层重要性采样。最后，我们使用 JAX JIT 编译、Adam 优化器、指数学习率衰减和梯度裁剪来训练模型，并使用 PSNR、深度与透明度可视化、采样诊断、360 度渲染及 marching-cubes 几何提取来评估新视角合成效果。
代码 · 101 行
import os, sys, subprocess, importlib.util, functools, dataclasses, time, math
def _sh(cmd):
subprocess.run(cmd, shell=True, check=False,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
print("Installing dependencies ...")
_sh(f'{sys.executable} -m pip install -q "etils[array-types,epy,etree,enp]" '
f'chex flax optax scikit-image')
REPO_DIR = "/content/jax3d" if os.path.isdir("/content") else os.path.abspath("./jax3d")
if not os.path.isdir(REPO_DIR):
print("Cloning google-research/jax3d ...")
_sh(f"git clone -q --depth 1 https://github.com/google-research/jax3d.git {REPO_DIR}")
def _load_module_by_path(name, path):
"""Load a single .py file without triggering the parent package __init__.
`from jax3d.math import volume_rendering` also works if you run
`pip install .` inside the clone, but that pulls in gin/tfds/etc.
"""
spec = importlib.util.spec_from_file_location(name, path)
mod = importlib.util.module_from_spec(spec)
sys.modules[name] = mod
spec.loader.exec_module(mod)
return mod
_VR_PATH = os.path.join(REPO_DIR, "jax3d", "jax3d", "math", "volume_rendering.py")
if not os.path.exists(_VR_PATH):
_VR_PATH = os.path.join(REPO_DIR, "jax3d", "math", "volume_rendering.py")
try:
j3vr = _load_module_by_path("j3d_volume_rendering", _VR_PATH)
except Exception as e:
raise SystemExit(
f"Could not load {_VR_PATH}: {e}\n"
"Try: pip install -U 'etils[array-types,epy,etree,enp]==1.9.4' and re-run."
)
import numpy as np
import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state
import matplotlib.pyplot as plt
from PIL import Image
print("jax", jax.__version__, "| device:", jax.devices()[0].device_kind,
f"({jax.devices()[0].platform})")
print("jax3d volume_rendering API:",
[n for n in ("sample_along_rays", "volume_rendering",
"sample_piecewise_constant_pdf", "sample_1d")
if hasattr(j3vr, n)])
@dataclasses.dataclass
class Config:
H: int = 64; W: int = 64
n_train_views: int = 24; n_test_views: int = 3
cam_radius: float = 3.2; fov_deg: float = 40.0
near: float = 1.9; far: float = 4.7
gt_samples: int = 256
n_coarse: int = 64; n_fine: int = 64
deg_pos: int = 10; deg_dir: int = 4
width: int = 128; depth: int = 6; skip: int = 3
batch_rays: int = 2048; steps: int = 2500
lr_init: float = 5e-4; lr_final: float = 5e-6
chunk: int = 4096
grid_res: int = 96
cfg = Config()
if jax.devices()[0].platform == "cpu":
print("\n!! No GPU detected -- switching to a small CPU-friendly config.")
print(" (Runtime > Change runtime type > T4 GPU for the full version.)\n")
cfg = dataclasses.replace(cfg, H=40, W=40, n_train_views=14, steps=400,
gt_samples=128, n_coarse=32, n_fine=32,
width=64, depth=4, skip=2, batch_rays=1024,
chunk=1600, grid_res=64)
def _normalize(v, axis=-1):
return v / (np.linalg.norm(v, axis=axis, keepdims=True) + 1e-9)
def look_at(eye, target=(0., 0., 0.), up=(0., 0., 1.)):
"""OpenGL/NeRF convention camera-to-world: +x right, +y up, camera looks at -z."""
eye, target, up = map(lambda a: np.asarray(a, np.float32), (eye, target, up))
fwd = _normalize(target - eye)
right = _normalize(np.cross(fwd, up))
trueup = np.cross(right, fwd)
c2w = np.eye(4, dtype=np.float32)
c2w[:3, :3] = np.stack([right, trueup, -fwd], axis=1)
c2w[:3, 3] = eye
return c2w
def orbit_poses(n, radius, elev_lo=18., elev_hi=58., phase=0.0):
"""Golden-angle azimuths + monotone elevations => well-spread views on a dome."""
i = np.arange(n, dtype=np.float64) + 0.5
az = 2 * np.pi * ((i * 0.6180339887) + phase)
elev = np.arcsin(np.linspace(np.sin(np.deg2rad(elev_lo)),
np.sin(np.deg2rad(elev_hi)), n))
eyes = np.stack([radius * np.cos(elev) * np.cos(az),
radius * np.cos(elev) * np.sin(az),
radius * np.sin(elev)], axis=-1).astype(np.float32)
return np.stack([look_at(e) for e in eyes], axis=0)
def rays_from_pose(c2w, H, W, focal):
"""Returns (origins, dirs) of shape [H, W, 3]; dirs are unit-length, so the
depths returned by jax3d's sampler are true world-space distances."""
i, j = np.meshgrid(np.arange(W, dtype=np.float32),
np.arange(H, dtype=np.float32), indexing="xy")
cam_dirs = np.stack([(i - W * .5 + .5) / focal,
-(j - H * .5 + .5) / focal,
-np.ones_like(i)], axis=-1)
dirs = _normalize(cam_dirs @ c2w[:3, :3].T)
origins = np.broadcast_to(c2w[:3, 3], dirs.shape)
return origins.astype(np.float32).copy(), dirs.astype(np.float32)
FOCAL = 0.5 * cfg.W / math.tan(0.5 * math.radians(cfg.fov_deg))代码 · 101 行
import os, sys, subprocess, importlib.util, functools, dataclasses, time, math
def _sh(cmd):
subprocess.run(cmd, shell=True, check=False,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
print("Installing dependencies ...")
_sh(f'{sys.executable} -m pip install -q "etils[array-types,epy,etree,enp]" '
f'chex flax optax scikit-image')
REPO_DIR = "/content/jax3d" if os.path.isdir("/content") else os.path.abspath("./jax3d")
if not os.path.isdir(REPO_DIR):
print("Cloning google-research/jax3d ...")
_sh(f"git clone -q --depth 1 https://github.com/google-research/jax3d.git {REPO_DIR}")
def _load_module_by_path(name, path):
"""Load a single .py file without triggering the parent package __init__.
`from jax3d.math import volume_rendering` also works if you run
`pip install .` inside the clone, but that pulls in gin/tfds/etc.
"""
spec = importlib.util.spec_from_file_location(name, path)
mod = importlib.util.module_from_spec(spec)
sys.modules[name] = mod
spec.loader.exec_module(mod)
return mod
_VR_PATH = os.path.join(REPO_DIR, "jax3d", "jax3d", "math", "volume_rendering.py")
if not os.path.exists(_VR_PATH):
_VR_PATH = os.path.join(REPO_DIR, "jax3d", "math", "volume_rendering.py")
try:
j3vr = _load_module_by_path("j3d_volume_rendering", _VR_PATH)
except Exception as e:
raise SystemExit(
f"Could not load {_VR_PATH}: {e}\n"
"Try: pip install -U 'etils[array-types,epy,etree,enp]==1.9.4' and re-run."
)
import numpy as np
import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state
import matplotlib.pyplot as plt
from PIL import Image
print("jax", jax.__version__, "| device:", jax.devices()[0].device_kind,
f"({jax.devices()[0].platform})")
print("jax3d volume_rendering API:",
[n for n in ("sample_along_rays", "volume_rendering",
"sample_piecewise_constant_pdf", "sample_1d")
if hasattr(j3vr, n)])
@dataclasses.dataclass
class Config:
H: int = 64; W: int = 64
n_train_views: int = 24; n_test_views: int = 3
cam_radius: float = 3.2; fov_deg: float = 40.0
near: float = 1.9; far: float = 4.7
gt_samples: int = 256
n_coarse: int = 64; n_fine: int = 64
deg_pos: int = 10; deg_dir: int = 4
width: int = 128; depth: int = 6; skip: int = 3
batch_rays: int = 2048; steps: int = 2500
lr_init: float = 5e-4; lr_final: float = 5e-6
chunk: int = 4096
grid_res: int = 96
cfg = Config()
if jax.devices()[0].platform == "cpu":
print("\n!! No GPU detected -- switching to a small CPU-friendly config.")
print(" (Runtime > Change runtime type > T4 GPU for the full version.)\n")
cfg = dataclasses.replace(cfg, H=40, W=40, n_train_views=14, steps=400,
gt_samples=128, n_coarse=32, n_fine=32,
width=64, depth=4, skip=2, batch_rays=1024,
chunk=1600, grid_res=64)
def _normalize(v, axis=-1):
return v / (np.linalg.norm(v, axis=axis, keepdims=True) + 1e-9)
def look_at(eye, target=(0., 0., 0.), up=(0., 0., 1.)):
"""OpenGL/NeRF convention camera-to-world: +x right, +y up, camera looks at -z."""
eye, target, up = map(lambda a: np.asarray(a, np.float32), (eye, target, up))
fwd = _normalize(target - eye)
right = _normalize(np.cross(fwd, up))
trueup = np.cross(right, fwd)
c2w = np.eye(4, dtype=np.float32)
c2w[:3, :3] = np.stack([right, trueup, -fwd], axis=1)
c2w[:3, 3] = eye
return c2w
def orbit_poses(n, radius, elev_lo=18., elev_hi=58., phase=0.0):
"""Golden-angle azimuths + monotone elevations => well-spread views on a dome."""
i = np.arange(n, dtype=np.float64) + 0.5
az = 2 * np.pi * ((i * 0.6180339887) + phase)
elev = np.arcsin(np.linspace(np.sin(np.deg2rad(elev_lo)),
np.sin(np.deg2rad(elev_hi)), n))
eyes = np.stack([radius * np.cos(elev) * np.cos(az),
radius * np.cos(elev) * np.sin(az),
radius * np.sin(elev)], axis=-1).astype(np.float32)
return np.stack([look_at(e) for e in eyes], axis=0)
def rays_from_pose(c2w, H, W, focal):
"""Returns (origins, dirs) of shape [H, W, 3]; dirs are unit-length, so the
depths returned by jax3d's sampler are true world-space distances."""
i, j = np.meshgrid(np.arange(W, dtype=np.float32),
np.arange(H, dtype=np.float32), indexing="xy")
cam_dirs = np.stack([(i - W * .5 + .5) / focal,
-(j - H * .5 + .5) / focal,
-np.ones_like(i)], axis=-1)
dirs = _normalize(cam_dirs @ c2w[:3, :3].T)
origins = np.broadcast_to(c2w[:3, 3], dirs.shape)
return origins.astype(np.float32).copy(), dirs.astype(np.float32)
FOCAL = 0.5 * cfg.W / math.tan(0.5 * math.radians(cfg.fov_deg))We set up the JAX3D environment, install the required dependencies, and load the volume_rendering module directly from the cloned repository. We configure GPU/CPU-adaptive training parameters and establish the camera model using pinhole intrinsics, look-at poses, and orbit-based camera placement. We then generate normalized world-space rays from each camera pose, providing the geometric foundation for the rendering pipeline.
我们配置 JAX3D 环境，安装所需的依赖项，并直接从克隆的仓库加载 volume_rendering 模块。我们设置适应 GPU/CPU 的训练参数，并使用针孔内参、look-at 姿态和基于轨道的相机放置来建立相机模型。随后，我们从每个相机姿态生成归一化的世界空间射线，为渲染管线提供几何基础。
更进一步：量化金融体系
看懂新闻只是起点——沿量化金融路径，把它变成能交付的工程能力