Radiance fields¶

inro.nerf, documented module by module; the radiance field tutorial shows the workflow.

Fields¶

fields ¶

Radiance fields assembled from spatial, color and direction modules.

As in Instant-NGP, a spatial network predicts density logits and geometry features, and a color network combines those features with an encoded view direction. Holding the color network and the renderer fixed compares spatial representations directly.

SphericalHarmonicsEncoding ¶

SphericalHarmonicsEncoding(degree: int = 4)

Real spherical-harmonics encoding of unit directions.

Evaluates the bands l = 0, ..., degree - 1, so degree counts bands, as in tiny-cuda-nn, and gives degree**2 features. degree=4 matches the view-direction encoding of Instant-NGP.

Parameters:

degree (int, default: 4 ) –

Number of bands, an integer from 1 to 4.

Shape

Input: (..., 3) unit directions.
Output: (..., degree**2).

Source code in inro/nerf/fields.py

def __init__(self, degree: int = 4) -> None:
    super().__init__()
    if not isinstance(degree, int) or not 1 <= degree <= 4:
        raise ValueError(f"Expected an integer 1 <= degree <= 4, got {degree!r}")
    self.degree = degree

out_features `property` ¶

out_features: int

Number of output features, degree**2.

IdentityEncoding ¶

Identity encoding of unit directions, as in FINER's NeRF network.

Shape

Input: (..., 3) unit directions.
Output: (..., 3).

RadianceField ¶

RadianceField(density_net: ModuleSpec, color_net: ModuleSpec, *, geometry_features: int = 15, direction_encoding: ModuleSpec = 'sh', aabb: Tensor | Sequence[float] = (-1.0, -1.0, -1.0, 1.0, 1.0, 1.0), density_activation: Callable[[Tensor], Tensor] = trunc_exp, density_bias: float = 0.0)

Radiance field with separate density, color and direction modules.

The density network maps points to one density logit followed by geometry features, and the color network maps geometry and direction features to RGB logits. This class applies the density activation and the RGB sigmoid, so neither network should have an output activation. Each module is a ModuleSpec: an instance, a registry key, a class or a (key_or_class, kwargs) pair.

Parameters:

density_net (ModuleSpec) –

Spatial network spec, without registry keys. A class receives in_features=3 and out_features=1 + geometry_features. Its inputs are coordinates in [-1, 1].
color_net (ModuleSpec) –

Color network spec, with registry keys "linear" and "mlp", returning three RGB logits. A class receives the input width and out_features=3.
geometry_features (int, default: 15 ) –

Number of features passed to the color network, excluding density; at least 1.
direction_encoding (ModuleSpec, default: 'sh' ) –

Direction encoding spec, exposing out_features, with registry keys "sh", degree-4 spherical harmonics as in Instant-NGP, and "identity", raw directions as in FINER.
aabb (Tensor | Sequence[float], default: (-1.0, -1.0, -1.0, 1.0, 1.0, 1.0) ) –

Scene box (xmin, ymin, zmin, xmax, ymax, zmax), kept in the state dict. World points map affinely to [-1, 1] and clamp at its boundary.
density_activation (Callable[[Tensor], Tensor], default: trunc_exp ) –

Callable from density logits to nonnegative densities.
density_bias (float, default: 0.0 ) –

Offset added to the density logits before the activation.

Shape

Input: x and unit directions, each (N, 3).
Output: (rgb, sigma), RGB (N, 3) in [0, 1] and density (N,).

Examples:

>>> import inro
>>> field = RadianceField(
...     (inro.InstantNGP, {"log2_hashmap_size": 14}),
...     ("mlp", {"hidden_features": 64, "hidden_layers": 2}),
... )
>>> directions = torch.tensor([0.0, 0.0, 1.0]).expand(8, 3)
>>> rgb, sigma = field(torch.rand(8, 3), directions)
>>> rgb.shape, sigma.shape
(torch.Size([8, 3]), torch.Size([8]))

Source code in inro/nerf/fields.py

def __init__(
    self,
    density_net: ModuleSpec,
    color_net: ModuleSpec,
    *,
    geometry_features: int = 15,
    direction_encoding: ModuleSpec = "sh",
    aabb: Tensor | Sequence[float] = (-1.0, -1.0, -1.0, 1.0, 1.0, 1.0),
    density_activation: Callable[[Tensor], Tensor] = trunc_exp,
    density_bias: float = 0.0,
) -> None:
    super().__init__()
    if geometry_features < 1:
        raise ValueError(f"geometry_features must be >= 1, got {geometry_features}")
    self.geometry_features = geometry_features
    self.density_net = build_module(
        density_net, {}, in_features=3, out_features=1 + geometry_features
    )
    self.register_buffer("aabb", torch.as_tensor(aabb, dtype=torch.float32))
    self.density_activation = density_activation
    self.density_bias = density_bias

    self.direction_encoding = build_module(direction_encoding, DIRECTION_ENCODINGS)
    color_in_features = geometry_features + self.direction_encoding.out_features
    self.color_net = build_module(
        color_net, COLOR_NETS, in_features=color_in_features, out_features=3
    )
    with torch.no_grad():
        # Supplied module instances may already use CUDA or half precision.
        parameter = next(self.color_net.parameters(), None)
        probe = torch.zeros(
            1,
            color_in_features,
            device=None if parameter is None else parameter.device,
            dtype=None if parameter is None else parameter.dtype,
        )
        out_features = int(self.color_net(probe).shape[-1])
    if out_features != 3:
        raise ValueError(
            f"color_net must map ({color_in_features},) -> (3,) RGB logits, but it "
            f"returns {out_features} features"
        )

normalize_coordinates ¶

normalize_coordinates(x: Tensor) -> Tensor

Map world coordinates from aabb to [-1, 1], clamping outliers.

Parameters:

x (Tensor) –

World coordinates (..., 3).

Source code in inro/nerf/fields.py

def normalize_coordinates(self, x: Tensor) -> Tensor:
    """Map world coordinates from :attr:`aabb` to ``[-1, 1]``, clamping outliers.

    Args:
        x: World coordinates ``(..., 3)``.
    """
    box_min, box_max = self.aabb[:3], self.aabb[3:]
    return (2.0 * (x - box_min) / (box_max - box_min) - 1.0).clamp(-1.0, 1.0)

query_density ¶

query_density(x: Tensor) -> Tensor

Return densities (N,), skipping the color network for occupancy updates.

Parameters:

x (Tensor) –

World coordinates (N, 3); the batch may be empty.

Source code in inro/nerf/fields.py

def query_density(self, x: Tensor) -> Tensor:
    """Return densities ``(N,)``, skipping the color network for occupancy updates.

    Args:
        x: World coordinates ``(N, 3)``; the batch may be empty.
    """
    if x.shape[0] == 0:
        return x.new_zeros((0,))
    return self._density_and_features(x)[0]

forward ¶

forward(x: Tensor, directions: Tensor) -> tuple[Tensor, Tensor]

Return RGB and density at world points seen along unit directions.

Parameters:

x (Tensor) –

World coordinates (N, 3); the batch may be empty.
directions (Tensor) –

Unit view directions (N, 3).

Source code in inro/nerf/fields.py

def forward(self, x: Tensor, directions: Tensor) -> tuple[Tensor, Tensor]:
    """Return RGB and density at world points seen along unit directions.

    Args:
        x: World coordinates ``(N, 3)``; the batch may be empty.
        directions: Unit view directions ``(N, 3)``.
    """
    if x.shape[0] == 0:
        return x.new_zeros((0, 3)), x.new_zeros((0,))
    sigma, features = self._density_and_features(x)
    direction_features = self.direction_encoding(directions)
    rgb = torch.sigmoid(
        self.color_net(torch.cat([features, direction_features], dim=-1))
    )
    return rgb, sigma

trunc_exp ¶

trunc_exp(x: Tensor) -> Tensor

Compute exp(x), clamping the input of its gradient to [-15, 15].

The clamp limits large density gradients early in training.

Parameters:

x (Tensor) –

Input of any shape.

Source code in inro/nerf/fields.py

def trunc_exp(x: Tensor) -> Tensor:
    """Compute ``exp(x)``, clamping the input of its gradient to ``[-15, 15]``.

    The clamp limits large density gradients early in training.

    Args:
        x: Input of any shape.
    """
    return _TruncExp.apply(x)

Renderers¶

renderers ¶

Occupancy-grid volume rendering with PyTorch and nerfacc backends.

Both backends march rays uniformly and composite emission and absorption. Ray directions must be unit vectors, so that near and far bounds, step size and depth are world distances. Call update_occupancy(field, step) before each training step.

Defaults follow torch-ngp's Blender recipe: the box [-1, 1]^3, near plane 0.2, a step of sqrt(3) * extent / 1024, and a 128**3 density grid thresholded at min(mean density, 10).

VolumeRenderer ¶

VolumeRenderer(aabb: Tensor | Sequence[float] = (-1.0, -1.0, -1.0, 1.0, 1.0, 1.0), near: float = 0.2, far: float | None = None, render_step_size: float | None = None, max_samples_per_ray: int | None = 1024, grid_resolution: int = 128, update_interval: int = 16, density_threshold: float = 10.0, ema_decay: float = 0.95, warmup_steps: int = 256)

Configuration and interface shared by occupancy-grid volume renderers.

Subclasses implement update_occupancy and forward. Build a backend with create_renderer.

Parameters:

aabb (Tensor | Sequence[float], default: (-1.0, -1.0, -1.0, 1.0, 1.0, 1.0) ) –

Scene box (xmin, ymin, zmin, xmax, ymax, zmax), left out of the state dict. Samples lie inside it.
near (float, default: 0.2 ) –

Near-plane distance; marching starts at the box entry if that is farther.
far (float | None, default: None ) –

Far-plane distance; None uses the box exit and stores infinity.
render_step_size (float | None, default: None ) –

Marching step; None uses sqrt(3) * extent / 1024, with extent the longest side of the box.
max_samples_per_ray (int | None, default: 1024 ) –

Most occupied samples per ray, front to back; None removes the cap. Samples past it count as empty. It bounds the queries of a radiance field, not the PyTorch backend's marching buffers.
grid_resolution (int, default: 128 ) –

Occupancy-grid cells per axis.
update_interval (int, default: 16 ) –

Training steps between occupancy updates.
density_threshold (float, default: 10.0 ) –

A cell is occupied when its density exceeds min(grid mean, density_threshold).
ema_decay (float, default: 0.95 ) –

Decay in the cell update max(ema_decay * density, sigma).
warmup_steps (int, default: 256 ) –

Steps during which an update covers every cell. Later updates sample num_cells // 4 uniform cells and up to that many occupied ones.

Attributes:

num_marching_steps –

Steps that cover the box diagonal, the longest crossing of a ray: ceil(diagonal / render_step_size), at least 1.

Source code in inro/nerf/renderers.py

def __init__(
    self,
    aabb: Tensor | Sequence[float] = (-1.0, -1.0, -1.0, 1.0, 1.0, 1.0),
    near: float = 0.2,
    far: float | None = None,
    render_step_size: float | None = None,
    max_samples_per_ray: int | None = 1024,
    grid_resolution: int = 128,
    update_interval: int = 16,
    density_threshold: float = 10.0,
    ema_decay: float = 0.95,
    warmup_steps: int = 256,
) -> None:
    super().__init__()
    aabb = torch.as_tensor(aabb, dtype=torch.float32)
    self.register_buffer("aabb", aabb, persistent=False)
    self.near = float(near)
    self.far = math.inf if far is None else float(far)
    if render_step_size is None:
        max_extent = (aabb[3:] - aabb[:3]).max().item()
        render_step_size = math.sqrt(3.0) * max_extent / 1024.0
    self.render_step_size = float(render_step_size)
    # The box diagonal bounds every ray's path through the scene.
    diagonal = (aabb[3:] - aabb[:3]).norm().item()
    self.num_marching_steps = max(math.ceil(diagonal / self.render_step_size), 1)
    self.max_samples_per_ray = (
        None if max_samples_per_ray is None else int(max_samples_per_ray)
    )
    self.grid_resolution = int(grid_resolution)
    self.update_interval = int(update_interval)
    self.density_threshold = float(density_threshold)
    self.ema_decay = float(ema_decay)
    self.warmup_steps = int(warmup_steps)

update_occupancy ¶

update_occupancy(field: Module, step: int) -> None

Update the occupancy grid from field.query_density.

Call once per training step, before rendering.

Parameters:

field (Module) –

Radiance field with the interface of RadianceField.
step (int) –

Zero-based training step. The grid updates only when step % update_interval == 0, all of it if step < warmup_steps.

Source code in inro/nerf/renderers.py

def update_occupancy(self, field: nn.Module, step: int) -> None:
    """Update the occupancy grid from ``field.query_density``.

    Call once per training step, before rendering.

    Args:
        field: Radiance field with the interface of :class:`RadianceField`.
        step: Zero-based training step. The grid updates only when
            ``step % update_interval == 0``, all of it if ``step < warmup_steps``.
    """
    raise NotImplementedError

forward ¶

forward(field: Module, rays_o: Tensor, rays_d: Tensor, background: Tensor | None = None) -> dict[str, Tensor | int]

Render a batch of rays through a radiance field.

Parameters:

field (Module) –

Radiance field with the interface of RadianceField.
rays_o (Tensor) –

Ray origins (N, 3).
rays_d (Tensor) –

Unit ray directions (N, 3).
background (Tensor | None, default: None ) –

Background color, shared (3,) or per ray (N, 3); None uses white.

Returns:

dict[str, Tensor | int] –
Dictionary with these keys.
- rgb: Colors (N, 3).
- opacity: Summed compositing weights (N, 1).
- depth: Expected distance along the ray given a hit, (N, 1).
- num_samples: Number of samples composited.

Source code in inro/nerf/renderers.py

def forward(
    self,
    field: nn.Module,
    rays_o: Tensor,
    rays_d: Tensor,
    background: Tensor | None = None,
) -> dict[str, Tensor | int]:
    """Render a batch of rays through a radiance field.

    Args:
        field: Radiance field with the interface of :class:`RadianceField`.
        rays_o: Ray origins ``(N, 3)``.
        rays_d: Unit ray directions ``(N, 3)``.
        background: Background color, shared ``(3,)`` or per ray ``(N, 3)``;
            ``None`` uses white.

    Returns:
        Dictionary with these keys.

        - ``rgb``: Colors ``(N, 3)``.
        - ``opacity``: Summed compositing weights ``(N, 1)``.
        - ``depth``: Expected distance along the ray given a hit, ``(N, 1)``.
        - ``num_samples``: Number of samples composited.
    """
    raise NotImplementedError

PyTorchRenderer ¶

PyTorchRenderer(**kwargs: Any)

Occupancy-grid volume renderer in pure PyTorch.

Marches each ray from box entry to exit at a fixed step and queries the radiance field at occupied samples only; skipped samples have zero density. The grid starts fully occupied, and it is in the state dict, as nerfacc's is, so a saved run renders with the grid it was scored with.

Parameters:

**kwargs (Any, default: {} ) –

Arguments of VolumeRenderer.

Attributes:

density_grid –

Running cell densities (G, G, G), for G = grid_resolution.
binaries –

Boolean occupancy mask (G, G, G), named after nerfacc's estimator.

Source code in inro/nerf/renderers.py

def __init__(self, **kwargs: Any) -> None:
    super().__init__(**kwargs)
    resolution = self.grid_resolution
    self.register_buffer(
        "density_grid", torch.zeros(resolution, resolution, resolution)
    )
    self.register_buffer(
        "binaries", torch.ones(resolution, resolution, resolution, dtype=torch.bool)
    )

update_occupancy ¶

update_occupancy(field: Module, step: int) -> None

Update selected cells from one jittered point each and threshold the grid.

Queries run in chunks of 2**16, in training or eval mode.

Parameters:

field (Module) –

Radiance field with the interface of RadianceField.
step (int) –

Zero-based training step. The grid updates only when step % update_interval == 0, all of it if step < warmup_steps.

Source code in inro/nerf/renderers.py

@torch.no_grad()
def update_occupancy(self, field: nn.Module, step: int) -> None:
    """Update selected cells from one jittered point each and threshold the grid.

    Queries run in chunks of ``2**16``, in training or eval mode.

    Args:
        field: Radiance field with the interface of :class:`RadianceField`.
        step: Zero-based training step. The grid updates only when
            ``step % update_interval == 0``, all of it if ``step < warmup_steps``.
    """
    if step % self.update_interval != 0:
        return
    resolution = self.grid_resolution
    device = self.density_grid.device
    num_cells = resolution**3
    if step < self.warmup_steps:
        cells = torch.arange(num_cells, device=device)
    else:
        # Discover new occupied regions while refreshing known ones.
        num_uniform = num_cells // 4
        uniform_cells = torch.randint(num_cells, (num_uniform,), device=device)
        occupied_cells = (
            self.binaries.reshape(-1).nonzero(as_tuple=False).squeeze(-1)
        )
        if occupied_cells.numel() > num_uniform:
            indices = torch.randint(
                occupied_cells.numel(), (num_uniform,), device=device
            )
            occupied_cells = occupied_cells[indices]
        cells = torch.cat([uniform_cells, occupied_cells])

    cell_indices = torch.stack(
        torch.unravel_index(cells, (resolution,) * 3), dim=-1
    )
    positions = (
        cell_indices.float() + torch.rand_like(cell_indices, dtype=torch.float32)
    ) / resolution
    box_min, box_max = self.aabb[:3], self.aabb[3:]
    positions = box_min + positions * (box_max - box_min)

    sampled_densities = torch.cat(
        [field.query_density(chunk) for chunk in positions.split(2**16)]
    )

    densities = self.density_grid.reshape(-1)
    densities[cells] = torch.maximum(
        densities[cells] * self.ema_decay, sampled_densities
    )
    threshold = min(self.density_threshold, densities.mean().item())
    self.binaries.copy_(self.density_grid > threshold)

forward ¶

forward(field: Module, rays_o: Tensor, rays_d: Tensor, background: Tensor | None = None) -> dict[str, Tensor | int]

Render a batch of rays through a radiance field.

Allocates (N, num_marching_steps) marching buffers. Each sample is jittered within its step in training mode and sits at its midpoint in eval mode.

Parameters:

field (Module) –

Radiance field with the interface of RadianceField.
rays_o (Tensor) –

Ray origins (N, 3).
rays_d (Tensor) –

Unit ray directions (N, 3).
background (Tensor | None, default: None ) –

Background color, shared (3,) or per ray (N, 3); None uses white.

Returns:

dict[str, Tensor | int] –

Dictionary with rgb, opacity, depth and num_samples; see VolumeRenderer.forward.

Source code in inro/nerf/renderers.py

def forward(
    self,
    field: nn.Module,
    rays_o: Tensor,
    rays_d: Tensor,
    background: Tensor | None = None,
) -> dict[str, Tensor | int]:
    """Render a batch of rays through a radiance field.

    Allocates ``(N, num_marching_steps)`` marching buffers. Each sample is jittered
    within its step in training mode and sits at its midpoint in eval mode.

    Args:
        field: Radiance field with the interface of :class:`RadianceField`.
        rays_o: Ray origins ``(N, 3)``.
        rays_d: Unit ray directions ``(N, 3)``.
        background: Background color, shared ``(3,)`` or per ray ``(N, 3)``;
            ``None`` uses white.

    Returns:
        Dictionary with ``rgb``, ``opacity``, ``depth`` and ``num_samples``; see
            :meth:`VolumeRenderer.forward`.
    """
    num_rays = rays_o.shape[0]
    device = rays_o.device
    step_size = self.render_step_size
    if background is None:
        background = torch.ones(3, device=device)
    background = background.to(device)

    # Starting at each box entry keeps the grid within the diagonal budget.
    entry_distances, exit_distances = self._ray_aabb_intersect(rays_o, rays_d)
    num_marching_steps = self.num_marching_steps
    steps = torch.arange(num_marching_steps, device=device, dtype=torch.float32)
    offsets = (
        torch.rand(num_rays, num_marching_steps, device=device)
        if self.training
        else torch.full((1, num_marching_steps), 0.5, device=device)
    )
    distances = entry_distances[:, None] + (steps + offsets) * step_size
    sample_mask = distances < exit_distances[:, None]

    positions = rays_o[:, None, :] + distances[..., None] * rays_d[:, None, :]
    box_min, box_max = self.aabb[:3], self.aabb[3:]
    # Fold one axis at a time to avoid an (N, S, 3) index buffer.
    resolution = self.grid_resolution
    flat_indices = None
    for axis in range(3):
        indices = (
            (
                (positions[..., axis] - box_min[axis])
                / (box_max[axis] - box_min[axis])
                * resolution
            )
            .int()
            .clamp_(0, resolution - 1)
        )
        # Flat indices can exceed int32 for large grids.
        flat_indices = (
            indices.long()
            if flat_indices is None
            else flat_indices * resolution + indices
        )
    sample_mask = sample_mask & self.binaries.reshape(-1)[flat_indices]
    del flat_indices
    if self.max_samples_per_ray is not None:
        # Keep the first occupied samples in front-to-back order.
        sample_mask = sample_mask & (
            sample_mask.cumsum(-1, dtype=torch.int32) <= self.max_samples_per_ray
        )

    sample_densities = torch.zeros(num_rays, num_marching_steps, device=device)
    sample_colors = torch.zeros(num_rays, num_marching_steps, 3, device=device)
    num_samples = int(sample_mask.sum().item())
    if num_samples > 0:
        points = positions[sample_mask]
        directions = rays_d[:, None, :].expand_as(positions)[sample_mask]
        # Free the (N, S, 3) buffer before querying the radiance field.
        del positions
        colors, densities = field(points, directions)
        sample_densities[sample_mask] = densities
        sample_colors[sample_mask] = colors

    alphas = 1.0 - torch.exp(-sample_densities * step_size)
    transmittance = torch.cumprod(
        torch.cat(
            [torch.ones(num_rays, 1, device=device), 1.0 - alphas + 1e-10], dim=-1
        ),
        dim=-1,
    )[:, :-1]
    weights = alphas * transmittance

    opacity = weights.sum(-1, keepdim=True)
    rgb = (weights[..., None] * sample_colors).sum(-2) + (
        1.0 - opacity
    ) * background
    # Condition depth on a hit, matching nerfacc's rendering convention.
    depth = (weights * distances).sum(-1, keepdim=True) / opacity.clamp_min(
        torch.finfo(rgb.dtype).eps
    )
    return {
        "rgb": rgb,
        "opacity": opacity,
        "depth": depth,
        "num_samples": num_samples,
    }

NerfAccRenderer ¶

NerfAccRenderer(**kwargs: Any)

Occupancy-grid volume renderer on nerfacc.OccGridEstimator and CUDA kernels.

Requires a working nerfacc extension; see is_nerfacc_available. Unlike PyTorchRenderer, it renders only background until update_occupancy runs, and sampling queries density to stop rays at transmittance below 1e-4, before the sample cap applies. The state dict holds the grid under estimator..

Parameters:

**kwargs (Any, default: {} ) –

Arguments of VolumeRenderer.

Source code in inro/nerf/renderers.py

def __init__(self, **kwargs: Any) -> None:
    super().__init__(**kwargs)
    if nerfacc is None:
        raise ImportError(
            "nerfacc is not installed; use backend='pytorch' or "
            "pip install nerfacc"
        )
    self.estimator = nerfacc.OccGridEstimator(
        roi_aabb=self.aabb.tolist(),
        resolution=self.grid_resolution,
        levels=1,
    )

update_occupancy ¶

update_occupancy(field: Module, step: int) -> None

Update the grid through OccGridEstimator.update_every_n_steps.

Queries all selected cells in one batch and requires training mode.

Parameters:

field (Module) –

Radiance field with the interface of RadianceField.
step (int) –

Zero-based training step. The grid updates only when step % update_interval == 0, all of it if step < warmup_steps.

Source code in inro/nerf/renderers.py

def update_occupancy(self, field: nn.Module, step: int) -> None:
    """Update the grid through ``OccGridEstimator.update_every_n_steps``.

    Queries all selected cells in one batch and requires training mode.

    Args:
        field: Radiance field with the interface of :class:`RadianceField`.
        step: Zero-based training step. The grid updates only when
            ``step % update_interval == 0``, all of it if ``step < warmup_steps``.
    """
    self.estimator.update_every_n_steps(
        step=step,
        occ_eval_fn=lambda x: field.query_density(x).unsqueeze(-1),
        occ_thre=self.density_threshold,
        ema_decay=self.ema_decay,
        warmup_steps=self.warmup_steps,
        n=self.update_interval,
    )

forward ¶

forward(field: Module, rays_o: Tensor, rays_d: Tensor, background: Tensor | None = None) -> dict[str, Tensor | int]

Render a batch of rays through a radiance field.

Samples are stratified in training mode.

Parameters:

field (Module) –

Radiance field with the interface of RadianceField.
rays_o (Tensor) –

Ray origins (N, 3).
rays_d (Tensor) –

Unit ray directions (N, 3).
background (Tensor | None, default: None ) –

Background color, shared (3,) or per ray (N, 3); None uses white.

Returns:

dict[str, Tensor | int] –

Dictionary with rgb, opacity, depth and num_samples; see VolumeRenderer.forward.

Source code in inro/nerf/renderers.py

def forward(
    self,
    field: nn.Module,
    rays_o: Tensor,
    rays_d: Tensor,
    background: Tensor | None = None,
) -> dict[str, Tensor | int]:
    """Render a batch of rays through a radiance field.

    Samples are stratified in training mode.

    Args:
        field: Radiance field with the interface of :class:`RadianceField`.
        rays_o: Ray origins ``(N, 3)``.
        rays_d: Unit ray directions ``(N, 3)``.
        background: Background color, shared ``(3,)`` or per ray ``(N, 3)``;
            ``None`` uses white.

    Returns:
        Dictionary with ``rgb``, ``opacity``, ``depth`` and ``num_samples``; see
            :meth:`VolumeRenderer.forward`.
    """
    num_rays = rays_o.shape[0]
    if background is None:
        background = torch.ones(3, device=rays_o.device)
    background = background.to(rays_o.device)

    def midpoints(t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor) -> Tensor:
        t_mid = (t_starts + t_ends)[:, None] / 2.0
        return rays_o[ray_indices] + rays_d[ray_indices] * t_mid

    def sigma_fn(t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor) -> Tensor:
        return field.query_density(midpoints(t_starts, t_ends, ray_indices))

    def rgb_sigma_fn(
        t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor
    ) -> tuple[Tensor, Tensor]:
        positions = midpoints(t_starts, t_ends, ray_indices)
        return field(positions, rays_d[ray_indices])

    ray_indices, t_starts, t_ends = self.estimator.sampling(
        rays_o,
        rays_d,
        sigma_fn=sigma_fn,
        near_plane=self.near,
        # nerfacc uses a finite far bound; box traversal still stops rays.
        far_plane=min(self.far, 1e10),
        render_step_size=self.render_step_size,
        stratified=self.training,
    )
    if self.max_samples_per_ray is not None and ray_indices.numel() > 0:
        # Samples are grouped by ray and ordered front to back.
        _, counts = torch.unique_consecutive(ray_indices, return_counts=True)
        segment_starts = torch.cumsum(counts, 0) - counts
        position_in_ray = torch.arange(
            ray_indices.shape[0], device=ray_indices.device
        ) - segment_starts.repeat_interleave(counts)
        keep = position_in_ray < self.max_samples_per_ray
        ray_indices = ray_indices[keep]
        t_starts = t_starts[keep]
        t_ends = t_ends[keep]
    rgb, opacity, depth, _ = nerfacc.rendering(
        t_starts,
        t_ends,
        ray_indices,
        n_rays=num_rays,
        rgb_sigma_fn=rgb_sigma_fn,
        render_bkgd=background,
    )
    return {
        "rgb": rgb,
        "opacity": opacity,
        "depth": depth,
        "num_samples": int(t_starts.shape[0]),
    }

is_nerfacc_available ¶

is_nerfacc_available() -> bool

Return whether nerfacc imports and its CUDA intersection kernel runs.

Some installations import but lack usable CUDA kernels, so this runs a small ray_aabb_intersect.

Source code in inro/nerf/renderers.py

def is_nerfacc_available() -> bool:
    """Return whether nerfacc imports and its CUDA intersection kernel runs.

    Some installations import but lack usable CUDA kernels, so this runs a small
    ``ray_aabb_intersect``.
    """
    if nerfacc is None or not torch.cuda.is_available():
        return False
    try:
        rays_o = torch.zeros(1, 3, device="cuda")
        rays_d = torch.tensor([[0.0, 0.0, 1.0]], device="cuda")
        aabb = torch.tensor([[-1.0, -1.0, -1.0, 1.0, 1.0, 1.0]], device="cuda")
        nerfacc.ray_aabb_intersect(rays_o, rays_d, aabb)
        return True
    except Exception:
        return False

create_renderer ¶

create_renderer(backend: Literal['auto', 'pytorch', 'nerfacc'] = 'auto', **kwargs: Any) -> VolumeRenderer

Create a volume renderer, preferring a working nerfacc backend.

Parameters:

backend (Literal['auto', 'pytorch', 'nerfacc'], default: 'auto' ) –

"nerfacc" for NerfAccRenderer, "pytorch" for PyTorchRenderer, or "auto" for nerfacc when is_nerfacc_available and PyTorch otherwise.
**kwargs (Any, default: {} ) –

Arguments of VolumeRenderer.

Source code in inro/nerf/renderers.py

def create_renderer(
    backend: Literal["auto", "pytorch", "nerfacc"] = "auto", **kwargs: Any
) -> "VolumeRenderer":
    """Create a volume renderer, preferring a working nerfacc backend.

    Args:
        backend: ``"nerfacc"`` for :class:`NerfAccRenderer`, ``"pytorch"`` for
            :class:`PyTorchRenderer`, or ``"auto"`` for nerfacc when
            :func:`is_nerfacc_available` and PyTorch otherwise.
        **kwargs: Arguments of :class:`VolumeRenderer`.
    """
    if backend == "auto":
        backend = "nerfacc" if is_nerfacc_available() else "pytorch"
    if backend == "nerfacc":
        return NerfAccRenderer(**kwargs)
    if backend == "pytorch":
        return PyTorchRenderer(**kwargs)
    raise ValueError(
        f"Unknown backend {backend!r}; use 'auto', 'nerfacc', or 'pytorch'"
    )

Datasets¶

datasets ¶

Posed RGBA images and ray sampling for NeRF synthetic (Blender) scenes.

Camera defaults follow torch-ngp's Blender recipe, as used by FINER: scene_scale=0.8, bound=1 and near=0.2. Images load at full resolution; downsample=4, skip=4 gives FINER's 25 training views at 200 x 200. Straight RGBA keeps the background a rendering choice.

BlenderDataset ¶

BlenderDataset(root: str | PathLike[str], split: str = 'train', downsample: int = 1, skip: int = 1, scene_scale: float = 0.8, bound: float = 1.0, near: float = 0.2, background: float | Sequence[float] | Tensor = 1.0)

Blender scene with full-view access and random ray sampling.

dataset[i] returns one view's rays and RGBA, and sample_rays generates rays from the stored poses, without a ray buffer. Poses keep Blender's convention, with translations multiplied by scene_scale.

Parameters:

root (str | PathLike[str]) –

Scene directory holding transforms_{split}.json and the images.
split (str, default: 'train' ) –

Split name, usually "train", "val" or "test".
downsample (int, default: 1 ) –

Integer factor of resolution reduction by area averaging; 4 gives 200 x 200 Blender images.
skip (int, default: 1 ) –

Keep every skip-th frame in file order; 4 selects 25 of the 100 training views, and evaluation uses 1.
scene_scale (float, default: 0.8 ) –

Scale applied to camera positions.
bound (float, default: 1.0 ) –

Half-extent of the scene box [-bound, bound]^3.
near (float, default: 0.2 ) –

Near-plane distance, stored for the renderer.
background (float | Sequence[float] | Tensor, default: 1.0 ) –

Background color in [0, 1], a scalar or three channels. White is the standard for evaluating Blender scenes.

Attributes:

images –

Straight RGBA images (views, H, W, 4) in [0, 1].
poses –

Camera-to-world poses (views, 4, 4), with scaled translations.
directions –

Camera-frame ray directions (H, W, 3), shared by all views.
radius –

Mean camera distance from the origin, after scaling.
focal –

Focal length in pixels at the loaded resolution.
height –

Image height H in pixels at the loaded resolution.
width –

Image width W in pixels at the loaded resolution.
far (float | None) –

Far-plane distance, always None: each ray ends at the box exit.
aabb –

Scene box (xmin, ymin, zmin, xmax, ymax, zmax), from bound.
background –

Background color (3,).

Source code in inro/nerf/datasets.py

def __init__(
    self,
    root: str | os.PathLike[str],
    split: str = "train",
    downsample: int = 1,
    skip: int = 1,
    scene_scale: float = 0.8,
    bound: float = 1.0,
    near: float = 0.2,
    background: float | Sequence[float] | Tensor = 1.0,
) -> None:
    super().__init__()
    self.root = Path(root)
    self.split = split
    self.downsample = int(downsample)
    self.skip = int(skip)
    self.scene_scale = float(scene_scale)
    self.bound = float(bound)
    self.near = float(near)
    self.far: float | None = None
    self.aabb = torch.tensor((-bound,) * 3 + (bound,) * 3, dtype=torch.float32)
    self.background = torch.as_tensor(background, dtype=torch.float32).expand(3)

    with open(self.root / f"transforms_{split}.json") as file:
        metadata = json.load(file)

    images: list[Tensor] = []
    poses: list[Tensor] = []
    for frame in metadata["frames"][:: self.skip]:
        path = self.root / f"{frame['file_path']}.png"
        with Image.open(path) as image:
            rgba = torch.from_numpy(np.asarray(image, dtype=np.float32) / 255.0)
        if rgba.shape[-1] == 3:
            rgba = torch.cat([rgba, torch.ones_like(rgba[..., :1])], dim=-1)
        elif rgba.shape[-1] != 4:
            raise ValueError(
                f"{path}: expected an RGB or RGBA image, got "
                f"{rgba.shape[-1]} channel(s)"
            )
        if self.downsample > 1:
            # Average straight RGBA to match the Blender training protocol.
            rgba = (
                F.interpolate(
                    rgba.permute(2, 0, 1).unsqueeze(0),
                    size=(
                        rgba.shape[0] // self.downsample,
                        rgba.shape[1] // self.downsample,
                    ),
                    mode="area",
                )
                .squeeze(0)
                .permute(1, 2, 0)
            )
        images.append(rgba)
        pose = torch.tensor(frame["transform_matrix"], dtype=torch.float32)
        pose[:3, 3] *= self.scene_scale
        poses.append(pose)
    self.images = torch.stack(images)
    self.poses = torch.stack(poses)

    self.radius = self.poses[:, :3, 3].norm(dim=-1).mean().item()
    self.height, self.width = self.images.shape[1:3]
    self.focal = focal_from_fov(float(metadata["camera_angle_x"]), self.width)
    self.directions = generate_ray_directions(self.height, self.width, self.focal)

to ¶

to(device: str | device) -> Self

Move the stored tensors to a device in place and return this dataset.

Parameters:

device (str | device) –

Device to move to; sampling then runs there, without transfers.

Source code in inro/nerf/datasets.py

def to(self, device: str | torch.device) -> Self:
    """Move the stored tensors to a device in place and return this dataset.

    Args:
        device: Device to move to; sampling then runs there, without transfers.
    """
    for name in ("images", "poses", "directions", "aabb", "background"):
        setattr(self, name, getattr(self, name).to(device))
    return self

composite ¶

composite(rgba: Tensor, background: Tensor | None = None) -> Tensor

Composite straight RGBA onto a background color.

Parameters:

rgba (Tensor) –

Straight RGBA (..., 4) in [0, 1].
background (Tensor | None, default: None ) –

Background color broadcastable to (..., 3); None uses background.

Returns:

Tensor –

Composited RGB (..., 3).

Source code in inro/nerf/datasets.py

def composite(self, rgba: Tensor, background: Tensor | None = None) -> Tensor:
    """Composite straight RGBA onto a background color.

    Args:
        rgba: Straight RGBA ``(..., 4)`` in ``[0, 1]``.
        background: Background color broadcastable to ``(..., 3)``; ``None`` uses
            :attr:`background`.

    Returns:
        Composited RGB ``(..., 3)``.
    """
    background = self.background if background is None else background
    rgb, alpha = rgba[..., :3], rgba[..., 3:]
    return rgb * alpha + background * (1.0 - alpha)

rays_for_pose ¶

rays_for_pose(c2w: Tensor) -> tuple[Tensor, Tensor]

Generate the rays of a full image for any camera pose.

Parameters:

c2w (Tensor) –

Camera-to-world pose (4, 4) or (3, 4), such as from pose_spherical, used as given: its translation must be scaled like poses.

Returns:

tuple[Tensor, Tensor] –

Ray origins and unit directions, each (H, W, 3), on the device of directions.

Source code in inro/nerf/datasets.py

def rays_for_pose(self, c2w: Tensor) -> tuple[Tensor, Tensor]:
    """Generate the rays of a full image for any camera pose.

    Args:
        c2w: Camera-to-world pose ``(4, 4)`` or ``(3, 4)``, such as from
            :func:`~inro.nerf.pose_spherical`, used as given: its translation must
            be scaled like :attr:`poses`.

    Returns:
        Ray origins and unit directions, each ``(H, W, 3)``, on the device of
            :attr:`directions`.
    """
    return generate_rays(c2w.to(self.directions.device), self.directions)

sample_rays ¶

sample_rays(num_rays: int, generator: Generator | None = None) -> dict[str, Tensor]

Sample rays uniformly over all views and pixels, with replacement.

Parameters:

num_rays (int) –

Number of rays N.
generator (Generator | None, default: None ) –

Random number generator, which may be on another device than the data, such as a CPU generator with data on CUDA. None uses the global generator of the data's device.

Returns:

dict[str, Tensor] –

Dictionary with rays_o and unit rays_d, each (N, 3), and rgba, straight RGBA (N, 4), on the device of images.

Source code in inro/nerf/datasets.py

def sample_rays(
    self, num_rays: int, generator: torch.Generator | None = None
) -> dict[str, Tensor]:
    """Sample rays uniformly over all views and pixels, with replacement.

    Args:
        num_rays: Number of rays ``N``.
        generator: Random number generator, which may be on another device than the
            data, such as a CPU generator with data on CUDA. ``None`` uses the
            global generator of the data's device.

    Returns:
        Dictionary with ``rays_o`` and unit ``rays_d``, each ``(N, 3)``, and
            ``rgba``, straight RGBA ``(N, 4)``, on the device of :attr:`images`.
    """
    device = self.images.device
    num_pixels = self.height * self.width
    # Draw on the generator's device to support CPU generators with CUDA data.
    generator_device = generator.device if generator is not None else device
    view_indices = torch.randint(
        len(self.images), (num_rays,), device=generator_device, generator=generator
    ).to(device)
    pixel_indices = torch.randint(
        num_pixels, (num_rays,), device=generator_device, generator=generator
    ).to(device)
    directions = self.directions.reshape(-1, 3)[pixel_indices]
    rays_o, rays_d = generate_rays(self.poses[view_indices], directions)
    rgba = self.images.reshape(len(self.images), -1, 4)[view_indices, pixel_indices]
    return {"rays_o": rays_o, "rays_d": rays_d, "rgba": rgba}

Training¶

training ¶

Per-scene NeRF training with random ray batches and mean squared error.

Adam keeps PyTorch's defaults, as in inro.train, so the learning rate is the one setting to tune. As in torch-ngp, it decays exponentially to 0.1 * lr and each step composites onto a fresh random background, while evaluation uses the dataset's. Results and logging follow inro.train.

train ¶

train(field: Module, renderer: VolumeRenderer, train_dataset: BlenderDataset, eval_dataset: BlenderDataset | None = None, *, num_steps: int = 25000, num_rays: int = 4096, max_samples_per_ray: int | None = 1024, lr: float = 0.01, color_lr: float | None = None, weight_decay: float = 0.0, adam_betas: tuple[float, float] = (0.9, 0.999), adam_eps: float = 1e-08, random_background: bool = True, metrics: Mapping[str, Callable[[Tensor, Tensor], float]] | None = None, eval_interval: int | None = None, eval_indices: Sequence[int] | None = None, log_interval: int = 100, device: str | device | None = None, log_dir: str | PathLike[str] | None = None, seed: int = 0) -> dict[str, Any]

Fit a radiance field to one scene with random ray batches.

Each step updates the occupancy grid, renders num_rays rays and minimizes their RGB mean squared error with Adam. A batch without occupied samples skips the optimizer update and counts towards skipped_steps.

Parameters:

field (Module) –

Radiance field, moved to device.
renderer (VolumeRenderer) –

Volume renderer, moved to device.
train_dataset (BlenderDataset) –

Posed images, moved to device in place.
eval_dataset (BlenderDataset | None, default: None ) –

Held-out views; None disables evaluation.
num_steps (int, default: 25000 ) –

Number of training steps.
num_rays (int, default: 4096 ) –

Rays sampled per step.
max_samples_per_ray (int | None, default: 1024 ) –

Most occupied samples per ray, front to back; None removes the cap. Overrides the renderer's setting.
lr (float, default: 0.01 ) –

Initial learning rate of Adam, decayed exponentially to 0.1 * lr.
color_lr (float | None, default: None ) –

Initial learning rate of field.color_net, decayed alike; None uses lr.
weight_decay (float, default: 0.0 ) –

Weight decay of Adam.
adam_betas (tuple[float, float], default: (0.9, 0.999) ) –

Moment decay rates of Adam.
adam_eps (float, default: 1e-08 ) –

Epsilon of Adam.
random_background (bool, default: True ) –

Whether to composite predictions and targets onto a fresh random color per ray, which discourages white haze in empty space. Otherwise, and always in evaluation, they use the dataset's background.
metrics (Mapping[str, Callable[[Tensor, Tensor], float]] | None, default: None ) –

Metrics passed to evaluate; None uses PSNR, SSIM and LPIPS.
eval_interval (int | None, default: None ) –

Steps between evaluations, the final step always included; None evaluates only there and non-positive values never.
eval_indices (Sequence[int] | None, default: None ) –

View indices to evaluate; None evaluates all views.
log_interval (int, default: 100 ) –

Steps between history records; final and evaluation steps are always recorded, and non-positive values record only those.
device (str | device | None, default: None ) –

Device to run on; None uses CUDA when available, else the CPU.
log_dir (str | PathLike[str] | None, default: None ) –

Directory for log.txt (truncated first), log.json (appended), the final checkpoint.pt of field and renderer.pt, the renderer's state with its occupancy grid; None writes no files.
seed (int, default: 0 ) –

Seed of the global PyTorch RNG, set after field and renderer are built, so it does not control their initialization.

Returns:

dict[str, Any] –
Dictionary with these keys.
- config: Settings of the run, with skipped_steps and peak_memory, the most bytes a training step allocated beyond what the device already held (zero off CUDA).
- model_dict: Final state dict of field.
- history: One record per logged step, with step, elapsed (seconds spent training so far), loss, psnr, num_samples and, if the step evaluated, the mean held-out metrics under eval.

Source code in inro/nerf/training.py

def train(
    field: nn.Module,
    renderer: VolumeRenderer,
    train_dataset: BlenderDataset,
    eval_dataset: BlenderDataset | None = None,
    *,
    num_steps: int = 25000,
    num_rays: int = 4096,
    max_samples_per_ray: int | None = 1024,
    lr: float = 1e-2,
    color_lr: float | None = None,
    weight_decay: float = 0.0,
    adam_betas: tuple[float, float] = (0.9, 0.999),
    adam_eps: float = 1e-8,
    random_background: bool = True,
    metrics: Mapping[str, Callable[[Tensor, Tensor], float]] | None = None,
    eval_interval: int | None = None,
    eval_indices: Sequence[int] | None = None,
    log_interval: int = 100,
    device: str | torch.device | None = None,
    log_dir: str | os.PathLike[str] | None = None,
    seed: int = 0,
) -> dict[str, Any]:
    """Fit a radiance field to one scene with random ray batches.

    Each step updates the occupancy grid, renders ``num_rays`` rays and minimizes their
    RGB mean squared error with Adam. A batch without occupied samples skips the
    optimizer update and counts towards ``skipped_steps``.

    Args:
        field: Radiance field, moved to ``device``.
        renderer: Volume renderer, moved to ``device``.
        train_dataset: Posed images, moved to ``device`` in place.
        eval_dataset: Held-out views; ``None`` disables evaluation.
        num_steps: Number of training steps.
        num_rays: Rays sampled per step.
        max_samples_per_ray: Most occupied samples per ray, front to back; ``None``
            removes the cap. Overrides the renderer's setting.
        lr: Initial learning rate of Adam, decayed exponentially to ``0.1 * lr``.
        color_lr: Initial learning rate of ``field.color_net``, decayed alike; ``None``
            uses ``lr``.
        weight_decay: Weight decay of Adam.
        adam_betas: Moment decay rates of Adam.
        adam_eps: Epsilon of Adam.
        random_background: Whether to composite predictions and targets onto a fresh
            random color per ray, which discourages white haze in empty space.
            Otherwise, and always in evaluation, they use the dataset's background.
        metrics: Metrics passed to :func:`evaluate`; ``None`` uses PSNR, SSIM and LPIPS.
        eval_interval: Steps between evaluations, the final step always included;
            ``None`` evaluates only there and non-positive values never.
        eval_indices: View indices to evaluate; ``None`` evaluates all views.
        log_interval: Steps between history records; final and evaluation steps are
            always recorded, and non-positive values record only those.
        device: Device to run on; ``None`` uses CUDA when available, else the CPU.
        log_dir: Directory for ``log.txt`` (truncated first), ``log.json`` (appended),
            the final ``checkpoint.pt`` of ``field`` and ``renderer.pt``, the renderer's
            state with its occupancy grid; ``None`` writes no files.
        seed: Seed of the global PyTorch RNG, set after ``field`` and ``renderer`` are
            built, so it does not control their initialization.

    Returns:
        Dictionary with these keys.

        - ``config``: Settings of the run, with ``skipped_steps`` and ``peak_memory``,
          the most bytes a training step allocated beyond what the device already held
          (zero off CUDA).
        - ``model_dict``: Final state dict of ``field``.
        - ``history``: One record per logged step, with ``step``, ``elapsed`` (seconds
          spent training so far), ``loss``, ``psnr``, ``num_samples`` and, if the step
          evaluated, the mean held-out metrics under ``eval``.
    """
    logger, json_logger = setup_experiment_loggers(log_dir, __name__)
    torch.manual_seed(seed)
    device = torch.device(device or ("cuda:0" if torch.cuda.is_available() else "cpu"))

    field = field.to(device)
    renderer = renderer.to(device)
    train_dataset.to(device)
    renderer.max_samples_per_ray = (
        None if max_samples_per_ray is None else int(max_samples_per_ray)
    )
    field.train()
    renderer.train()

    num_parameters = count_parameters(field)
    config = {
        "model": type(field).__name__,
        "density_net": type(getattr(field, "density_net", None)).__name__,
        "renderer": type(renderer).__name__,
        "num_params": num_parameters,
        "num_steps": num_steps,
        "num_rays": num_rays,
        "max_samples_per_ray": max_samples_per_ray,
        "lr": lr,
        "color_lr": color_lr,
        "weight_decay": weight_decay,
        "adam_betas": tuple(adam_betas),
        "adam_eps": adam_eps,
        "random_background": random_background,
        "seed": seed,
        "device": str(device),
        "num_train_views": len(train_dataset),
        "image_size": (train_dataset.height, train_dataset.width),
        "aabb": renderer.aabb.tolist(),
        "near": renderer.near,
        "far": renderer.far,
        "render_step_size": renderer.render_step_size,
        "grid_resolution": renderer.grid_resolution,
        "density_threshold": renderer.density_threshold,
    }
    logger.debug("Config: %s", config)

    groups = [{"params": field.parameters()}]
    if color_lr is not None:
        named = field.named_parameters()
        groups = [
            {"params": [p for n, p in named if not n.startswith("color_net.")]},
            {"params": field.color_net.parameters(), "lr": color_lr},
        ]
    optimizer = torch.optim.Adam(
        groups,
        lr=lr,
        betas=adam_betas,
        eps=adam_eps,
        weight_decay=weight_decay,
    )
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lambda step: 0.1 ** min(step / num_steps, 1.0)
    )

    if eval_interval is None:
        eval_interval = num_steps
    history: list[dict[str, Any]] = []
    status = ""
    elapsed = 0.0
    skipped_steps = 0
    progress = tqdm(range(1, num_steps + 1), desc="Training")
    cuda = device.type == "cuda"
    resident = torch.cuda.memory_allocated(device) if cuda else 0
    peak_memory = 0

    start_time = time.perf_counter()
    for step in progress:
        if cuda:  # each step's own peak, without evaluations
            torch.cuda.reset_peak_memory_stats(device)

        # Step zero fills nerfacc's initially empty occupancy grid.
        renderer.update_occupancy(field, step - 1)

        batch = train_dataset.sample_rays(num_rays)
        rays_o, rays_d, rgba = batch["rays_o"], batch["rays_d"], batch["rgba"]

        background = (
            torch.rand_like(rgba[..., :3])
            if random_background
            else train_dataset.background
        )
        target = train_dataset.composite(rgba, background)

        result = renderer(field, rays_o, rays_d, background=background)
        loss = F.mse_loss(result["rgb"], target)

        optimizer.zero_grad()
        # Pure-background batches do not depend on the parameters.
        if loss.requires_grad:
            loss.backward()
            optimizer.step()
        else:
            skipped_steps += 1
        scheduler.step()

        if cuda:
            peak = torch.cuda.max_memory_allocated(device) - resident
            peak_memory = max(peak_memory, peak)

        should_log = (log_interval > 0 and step % log_interval == 0) or (
            step == num_steps
        )
        should_eval = (
            eval_dataset is not None
            and eval_interval > 0
            and (step % eval_interval == 0 or step == num_steps)
        )
        if not (should_log or should_eval):
            continue  # no host sync, so the GPU never waits for Python

        if cuda:
            # CUDA kernels run asynchronously; wait so timings are wall-clock.
            torch.cuda.synchronize(device)
        elapsed += time.perf_counter() - start_time
        mse = loss.detach().item()
        entry: dict[str, Any] = {
            "step": step,
            "elapsed": elapsed,
            "loss": mse,
            "psnr": 10 * math.log10(1 / max(mse, 1e-12)),
            "num_samples": result["num_samples"],
        }
        status = _format_metrics(entry)
        progress.set_postfix_str(status)

        if should_eval:
            entry["eval"] = evaluate(
                field,
                renderer,
                eval_dataset,
                indices=eval_indices,
                metrics=metrics,
            )["mean"]
            status += " - eval " + _format_metrics(entry["eval"])

        logger.info("Step [%d/%d] - %s", step, num_steps, status)
        json_logger.log(entry)
        history.append(entry)
        start_time = time.perf_counter()  # evaluating and logging is not training

    logger.info("Training completed in %.1f s - %s", elapsed, status)
    if skipped_steps:
        warnings.warn(
            f"{skipped_steps} of {num_steps} steps had no occupied samples and "
            f"were skipped; the field's occupancy grid is empty or nearly so. "
            f"Treat this run's metrics as invalid rather than as a score.",
            stacklevel=2,
        )
    config["skipped_steps"] = skipped_steps
    config["peak_memory"] = peak_memory

    if log_dir:
        checkpoint_path = os.path.join(log_dir, "checkpoint.pt")
        torch.save(field.state_dict(), checkpoint_path)
        # The grid decides which samples a render queries, so rendering the saved
        # radiance field needs the grid the run was scored with.
        torch.save(renderer.state_dict(), os.path.join(log_dir, "renderer.pt"))
        logger.info("Saved checkpoint to %s", checkpoint_path)

    return {"config": config, "model_dict": field.state_dict(), "history": history}

Evaluation and rendering¶

evaluation ¶

Rendering and held-out-view evaluation for radiance fields.

render_image ¶

render_image(field: Module, renderer: VolumeRenderer, rays_o: Tensor, rays_d: Tensor, *, background: Tensor | None = None, chunk_size: int | None = None) -> dict[str, Tensor]

Render rays in chunks, without gradients or stratified jitter.

Both modules enter eval mode, and return to training mode if field was training.

Parameters:

field (Module) –

Radiance field; rendering runs on the device of its parameters.
renderer (VolumeRenderer) –

Volume renderer on the device of field, such as from create_renderer.
rays_o (Tensor) –

Ray origins (..., 3), such as (H, W, 3) for a view.
rays_d (Tensor) –

Unit ray directions (..., 3).
background (Tensor | None, default: None ) –

Background color (3,); None uses white.
chunk_size (int | None, default: None ) –

Rays per forward pass; None fits 2**19 marching samples.

Returns:

dict[str, Tensor] –

Dictionary with rgb (..., 3), opacity (..., 1) and depth (..., 1), with the leading dimensions of rays_o.

Source code in inro/nerf/evaluation.py

@torch.no_grad()
def render_image(
    field: nn.Module,
    renderer: VolumeRenderer,
    rays_o: Tensor,
    rays_d: Tensor,
    *,
    background: Tensor | None = None,
    chunk_size: int | None = None,
) -> dict[str, Tensor]:
    """Render rays in chunks, without gradients or stratified jitter.

    Both modules enter eval mode, and return to training mode if ``field`` was training.

    Args:
        field: Radiance field; rendering runs on the device of its parameters.
        renderer: Volume renderer on the device of ``field``, such as from
            :func:`create_renderer`.
        rays_o: Ray origins ``(..., 3)``, such as ``(H, W, 3)`` for a view.
        rays_d: Unit ray directions ``(..., 3)``.
        background: Background color ``(3,)``; ``None`` uses white.
        chunk_size: Rays per forward pass; ``None`` fits ``2**19`` marching samples.

    Returns:
        Dictionary with ``rgb`` ``(..., 3)``, ``opacity`` ``(..., 1)`` and ``depth``
            ``(..., 1)``, with the leading dimensions of ``rays_o``.
    """
    if chunk_size is None:
        # Marching buffers are allocated before the per-ray sample cap is applied.
        chunk_size = max(2**19 // getattr(renderer, "num_marching_steps", 1024), 1)
    ray_shape = rays_o.shape[:-1]
    device = next(field.parameters()).device
    rays_o = rays_o.reshape(-1, 3).to(device)
    rays_d = rays_d.reshape(-1, 3).to(device)
    outputs: dict[str, list[Tensor]] = {"rgb": [], "opacity": [], "depth": []}
    with _eval_mode(field, renderer):
        for start in range(0, rays_o.shape[0], chunk_size):
            result = renderer(
                field,
                rays_o[start : start + chunk_size],
                rays_d[start : start + chunk_size],
                background=background,
            )
            for key in outputs:
                outputs[key].append(result[key])
    return {
        key: torch.cat(values).reshape(*ray_shape, -1)
        for key, values in outputs.items()
    }

evaluate ¶

evaluate(field: Module, renderer: VolumeRenderer, dataset: BlenderDataset, *, indices: Iterable[int] | None = None, metrics: Mapping[str, Callable[[Tensor, Tensor], float]] | None = None, chunk_size: int | None = None, return_images: bool = False) -> dict[str, Any]

Score held-out views, composited onto the dataset's background.

Renders are clamped to [0, 1] before scoring. Both modules enter eval mode, and return to training mode if field was training.

Parameters:

field (Module) –

Radiance field; rendering runs on the device of its parameters.
renderer (VolumeRenderer) –

Volume renderer on the device of field, such as from create_renderer.
dataset (BlenderDataset) –

Posed RGBA images and the background they are composited onto.
indices (Iterable[int] | None, default: None ) –

View indices; None evaluates all views.
metrics (Mapping[str, Callable[[Tensor, Tensor], float]] | None, default: None ) –

Metrics by name, called as fn(pred, target) on images (3, H, W) in [0, 1]; None uses PSNR, SSIM and LPIPS (VGG).
chunk_size (int | None, default: None ) –

Rays per forward pass; None fits 2**19 marching samples.
return_images (bool, default: False ) –

Whether to include the rendered images in the result.

Returns:

dict[str, Any] –
Dictionary with these keys.
- views: One record per view, with view (its index), duration (seconds of its render), peak_memory (the most bytes it allocated beyond what the device already held, zero off CUDA) and the metric values.
- mean: Average of these entries except the index, empty without views.
- images: Rendered CPU tensors (H, W, 3), if return_images.

Source code in inro/nerf/evaluation.py

@torch.no_grad()
def evaluate(
    field: nn.Module,
    renderer: VolumeRenderer,
    dataset: BlenderDataset,
    *,
    indices: Iterable[int] | None = None,
    metrics: Mapping[str, Callable[[Tensor, Tensor], float]] | None = None,
    chunk_size: int | None = None,
    return_images: bool = False,
) -> dict[str, Any]:
    """Score held-out views, composited onto the dataset's background.

    Renders are clamped to ``[0, 1]`` before scoring. Both modules enter eval mode, and
    return to training mode if ``field`` was training.

    Args:
        field: Radiance field; rendering runs on the device of its parameters.
        renderer: Volume renderer on the device of ``field``, such as from
            :func:`create_renderer`.
        dataset: Posed RGBA images and the background they are composited onto.
        indices: View indices; ``None`` evaluates all views.
        metrics: Metrics by name, called as ``fn(pred, target)`` on images ``(3, H, W)``
            in ``[0, 1]``; ``None`` uses PSNR, SSIM and LPIPS (VGG).
        chunk_size: Rays per forward pass; ``None`` fits ``2**19`` marching samples.
        return_images: Whether to include the rendered images in the result.

    Returns:
        Dictionary with these keys.

        - ``views``: One record per view, with ``view`` (its index), ``duration``
          (seconds of its render), ``peak_memory`` (the most bytes it allocated beyond
          what the device already held, zero off CUDA) and the metric values.
        - ``mean``: Average of these entries except the index, empty without views.
        - ``images``: Rendered CPU tensors ``(H, W, 3)``, if ``return_images``.
    """
    if metrics is None:
        metrics = {"psnr": psnr, "ssim": ssim, "lpips": lpips}
    indices = range(len(dataset)) if indices is None else indices
    device = next(field.parameters()).device
    background = dataset.background.to(device)

    cuda = device.type == "cuda"
    views: list[dict[str, int | float]] = []
    images: list[Tensor] = []
    with _eval_mode(field, renderer):
        for index in indices:
            batch = dataset[index]
            if cuda:  # the view's own time and peak, after any work queued before
                torch.cuda.synchronize(device)
                torch.cuda.reset_peak_memory_stats(device)
            resident = torch.cuda.memory_allocated(device) if cuda else 0
            start_time = time.time()
            rendered = render_image(
                field,
                renderer,
                batch["rays_o"],
                batch["rays_d"],
                background=background,
                chunk_size=chunk_size,
            )
            if cuda:
                torch.cuda.synchronize(device)
            duration = time.time() - start_time
            peak = torch.cuda.max_memory_allocated(device) - resident if cuda else 0
            pred = rendered["rgb"].clamp(0.0, 1.0)
            target = dataset.composite(batch["rgba"].to(device), background)
            entry: dict[str, int | float] = {
                "view": int(index),
                "duration": duration,
                "peak_memory": peak,
            }
            for name, metric in metrics.items():
                entry[name] = metric(pred.permute(2, 0, 1), target.permute(2, 0, 1))
            views.append(entry)
            if return_images:
                images.append(pred.cpu())

    keys = [key for key in views[0] if key != "view"] if views else []
    mean = {key: sum(view[key] for view in views) / len(views) for key in keys}
    result: dict[str, Any] = {"views": views, "mean": mean}
    if return_images:
        result["images"] = images
    return result

render_views ¶

render_views(field: Module, renderer: VolumeRenderer, dataset: BlenderDataset, poses: Sequence[Tensor], *, chunk_size: int | None = None) -> list[NDArray[uint8]]

Render camera poses, such as an orbit, as uint8 frames for a video.

Both modules enter eval mode, and return to training mode if field was training.

Parameters:

field (Module) –

Radiance field; rendering runs on the device of its parameters.
renderer (VolumeRenderer) –

Volume renderer on the device of field, such as from create_renderer.
dataset (BlenderDataset) –

Source of the camera intrinsics and the background color.
poses (Sequence[Tensor]) –

Camera-to-world poses (4, 4), such as from pose_spherical.
chunk_size (int | None, default: None ) –

Rays per forward pass; None fits 2**19 marching samples.

Returns:

list[NDArray[uint8]] –

List of uint8 NumPy arrays (H, W, 3), one per pose.

Source code in inro/nerf/evaluation.py

@torch.no_grad()
def render_views(
    field: nn.Module,
    renderer: VolumeRenderer,
    dataset: BlenderDataset,
    poses: Sequence[Tensor],
    *,
    chunk_size: int | None = None,
) -> list[NDArray[np.uint8]]:
    """Render camera poses, such as an orbit, as ``uint8`` frames for a video.

    Both modules enter eval mode, and return to training mode if ``field`` was training.

    Args:
        field: Radiance field; rendering runs on the device of its parameters.
        renderer: Volume renderer on the device of ``field``, such as from
            :func:`create_renderer`.
        dataset: Source of the camera intrinsics and the background color.
        poses: Camera-to-world poses ``(4, 4)``, such as from :func:`pose_spherical`.
        chunk_size: Rays per forward pass; ``None`` fits ``2**19`` marching samples.

    Returns:
        List of ``uint8`` NumPy arrays ``(H, W, 3)``, one per pose.
    """
    background = dataset.background.to(next(field.parameters()).device)
    frames: list[NDArray[np.uint8]] = []
    with _eval_mode(field, renderer):
        for pose in poses:
            rays_o, rays_d = dataset.rays_for_pose(pose)
            rendered = render_image(
                field,
                renderer,
                rays_o,
                rays_d,
                background=background,
                chunk_size=chunk_size,
            )
            rgb = (rendered["rgb"].clamp(0, 1) * 255).round().to(torch.uint8)
            frames.append(rgb.cpu().numpy())
    return frames

Cameras¶

cameras ¶

Camera utilities for novel view synthesis.

Follows the Blender and OpenGL convention of the NeRF synthetic dataset: in the camera frame, x points right, y up, and the camera looks along -z. A camera-to-world pose is a (4, 4) matrix whose upper (3, 4) block is [R | t].

focal_from_fov ¶

focal_from_fov(fov: float, size: int) -> float

Convert a full field of view to a focal length in pixels.

Parameters:

fov (float) –

Full field of view in radians, such as camera_angle_x of the Blender transforms files, which pairs with the image width.
size (int) –

Image size in pixels along the axis of fov.

Source code in inro/nerf/cameras.py

def focal_from_fov(fov: float, size: int) -> float:
    """Convert a full field of view to a focal length in pixels.

    Args:
        fov: Full field of view in radians, such as ``camera_angle_x`` of the Blender
            transforms files, which pairs with the image width.
        size: Image size in pixels along the axis of ``fov``.
    """
    return 0.5 * size / math.tan(0.5 * fov)

generate_ray_directions ¶

generate_ray_directions(height: int, width: int, focal: float, cx: float | None = None, cy: float | None = None, *, device: str | device | None = None, dtype: dtype = float32) -> Tensor

Generate camera-frame ray directions (H, W, 3) through the pixel centers.

Directions have z = -1. They stay unnormalized, so that poses with the same intrinsics can share them.

Parameters:

height (int) –

Image height H in pixels.
width (int) –

Image width W in pixels.
focal (float) –

Focal length in pixels.
cx (float | None, default: None ) –

Principal point x in pixels; None uses width / 2.
cy (float | None, default: None ) –

Principal point y in pixels; None uses height / 2.
device (str | device | None, default: None ) –

Device to run on; None uses the default device.
dtype (dtype, default: float32 ) –

Floating-point dtype of the directions.

Source code in inro/nerf/cameras.py

def generate_ray_directions(
    height: int,
    width: int,
    focal: float,
    cx: float | None = None,
    cy: float | None = None,
    *,
    device: str | torch.device | None = None,
    dtype: torch.dtype = torch.float32,
) -> Tensor:
    """Generate camera-frame ray directions ``(H, W, 3)`` through the pixel centers.

    Directions have ``z = -1``. They stay unnormalized, so that poses with the same
    intrinsics can share them.

    Args:
        height: Image height ``H`` in pixels.
        width: Image width ``W`` in pixels.
        focal: Focal length in pixels.
        cx: Principal point x in pixels; ``None`` uses ``width / 2``.
        cy: Principal point y in pixels; ``None`` uses ``height / 2``.
        device: Device to run on; ``None`` uses the default device.
        dtype: Floating-point dtype of the directions.
    """
    cx = 0.5 * width if cx is None else cx
    cy = 0.5 * height if cy is None else cy
    pixel_y, pixel_x = torch.meshgrid(
        torch.arange(height, device=device, dtype=dtype) + 0.5,
        torch.arange(width, device=device, dtype=dtype) + 0.5,
        indexing="ij",
    )
    return torch.stack(
        [(pixel_x - cx) / focal, -(pixel_y - cy) / focal, -torch.ones_like(pixel_x)],
        dim=-1,
    )

generate_rays ¶

generate_rays(c2w: Tensor, directions: Tensor, normalize: bool = True) -> tuple[Tensor, Tensor]

Generate world-space rays from camera poses and camera-frame directions.

Parameters:

c2w (Tensor) –

Camera-to-world poses (..., 4, 4) or (..., 3, 4).
directions (Tensor) –

Camera-frame directions (..., 3), whose leading dimensions broadcast with those of c2w: one pose (4, 4) with a grid (H, W, 3), or poses (N, 4, 4) with directions (N, 3).
normalize (bool, default: True ) –

Whether to normalize the directions, so that distances along a ray are in world units.

Returns:

tuple[Tensor, Tensor] –

Ray origins and directions, each (..., 3). The origins are an expanded view of the translations of c2w.

Source code in inro/nerf/cameras.py

def generate_rays(
    c2w: Tensor, directions: Tensor, normalize: bool = True
) -> tuple[Tensor, Tensor]:
    """Generate world-space rays from camera poses and camera-frame directions.

    Args:
        c2w: Camera-to-world poses ``(..., 4, 4)`` or ``(..., 3, 4)``.
        directions: Camera-frame directions ``(..., 3)``, whose leading dimensions
            broadcast with those of ``c2w``: one pose ``(4, 4)`` with a grid
            ``(H, W, 3)``, or poses ``(N, 4, 4)`` with directions ``(N, 3)``.
        normalize: Whether to normalize the directions, so that distances along a ray
            are in world units.

    Returns:
        Ray origins and directions, each ``(..., 3)``. The origins are an expanded view
            of the translations of ``c2w``.
    """
    rotation = c2w[..., :3, :3]
    translation = c2w[..., :3, 3]
    # Elementwise multiplication supports a different pose for each ray.
    rays_d = (rotation * directions[..., None, :]).sum(-1)
    rays_o = translation.expand_as(rays_d)
    if normalize:
        rays_d = F.normalize(rays_d, dim=-1)
    return rays_o, rays_d

pose_spherical ¶

pose_spherical(theta: float, phi: float, radius: float) -> Tensor

Compute a camera-to-world pose on a sphere, looking at the origin.

Uses the original NeRF convention, with world +z up. At theta = phi = 0 the camera sits at (0, radius, 0).

Parameters:

theta (float) –

Azimuth in degrees from +y, clockwise about +z seen from above: 90 places the camera on +x.
phi (float) –

Pitch in degrees. Negative values place the camera above the xy plane; the original NeRF render path uses -30.
radius (float) –

Distance from the origin.

Returns:

Tensor –

Pose (4, 4) on the default device, with the default dtype.

Source code in inro/nerf/cameras.py

def pose_spherical(theta: float, phi: float, radius: float) -> Tensor:
    """Compute a camera-to-world pose on a sphere, looking at the origin.

    Uses the original NeRF convention, with world ``+z`` up. At ``theta = phi = 0`` the
    camera sits at ``(0, radius, 0)``.

    Args:
        theta: Azimuth in degrees from ``+y``, clockwise about ``+z`` seen from above:
            ``90`` places the camera on ``+x``.
        phi: Pitch in degrees. Negative values place the camera above the ``xy`` plane;
            the original NeRF render path uses ``-30``.
        radius: Distance from the origin.

    Returns:
        Pose ``(4, 4)`` on the default device, with the default dtype.
    """
    c2w = torch.eye(4)
    c2w[2, 3] = radius

    phi = math.radians(phi)
    pitch = torch.eye(4)
    pitch[1, 1] = math.cos(phi)
    pitch[1, 2] = -math.sin(phi)
    pitch[2, 1] = math.sin(phi)
    pitch[2, 2] = math.cos(phi)
    c2w = pitch @ c2w

    theta = math.radians(theta)
    yaw = torch.eye(4)
    yaw[0, 0] = math.cos(theta)
    yaw[0, 2] = -math.sin(theta)
    yaw[2, 0] = math.sin(theta)
    yaw[2, 2] = math.cos(theta)
    c2w = yaw @ c2w

    # Convert the orbit's rotation axes to Blender's world frame.
    blender_axes = torch.tensor(
        [[-1.0, 0, 0, 0], [0, 0, 1.0, 0], [0, 1.0, 0, 0], [0, 0, 0, 1.0]]
    )
    return blender_axes @ c2w

Radiance fields¶

Fields¶

fields ¶

SphericalHarmonicsEncoding ¶

out_features property ¶

IdentityEncoding ¶

RadianceField ¶

normalize_coordinates ¶

query_density ¶

forward ¶

trunc_exp ¶

Renderers¶

renderers ¶

VolumeRenderer ¶

update_occupancy ¶

forward ¶

PyTorchRenderer ¶

update_occupancy ¶

forward ¶

NerfAccRenderer ¶

update_occupancy ¶

forward ¶

is_nerfacc_available ¶

create_renderer ¶

Datasets¶

datasets ¶

BlenderDataset ¶

to ¶

composite ¶

rays_for_pose ¶

sample_rays ¶

Training¶

training ¶

train ¶

Evaluation and rendering¶

evaluation ¶

render_image ¶

evaluate ¶

render_views ¶

Cameras¶

cameras ¶

focal_from_fov ¶

generate_ray_directions ¶

generate_rays ¶

pose_spherical ¶

out_features `property` ¶