NeRF volume rendering in PyTorch

The C++ skill volume rendering for NeRF derived the rendering sum and wrote it one ray at a time, with a loop and a running transmittance. Research code never looks like that. Open run_nerf.py in nerf-pytorch and you find the same mathematics written for thousands of rays at once, without a single loop over rays or samples, differentiable end to end so that the photo loss can train the MLP behind it. This lesson is about writing it that way: the shapes, the one-line tricks (an exclusive product with cumprod, inverse-CDF sampling with searchsorted), and the handful of bugs that every NeRF re-implementation hits once.

Everything here is the code around the network — the part that is the same in NeRF, mip-NeRF, Instant-NGP's PyTorch ports and nerfstudio's samplers and renderers. Reference points throughout:

1. Shapes first

Every tensor in a NeRF renderer has the shape (*batch, N, C) or a prefix of it: *batch for the rays (usually one axis, R), N for the samples along a ray, C for channels (3 for points and colours, 4 for raw outputs). The whole skill is a discipline of writing every operation so that it works for any *batch — which means you never index the batch axes. Reductions over samples say dim=-1 (or -2 when a channel axis follows), and per-ray values are broadcast back with [..., None].

The first line that uses this is the one that turns depths into points:

import torch

R, N = 4, 6
rays_o = torch.zeros(R, 3)
rays_d = torch.tensor([[0.0, 0.0, -1.0]]).expand(R, 3)
z_vals = torch.linspace(2.0, 6.0, N).expand(R, N)

pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals[..., :, None]
print("rays_d[..., None, :]", tuple(rays_d[..., None, :].shape))
print("z_vals[..., :, None]", tuple(z_vals[..., :, None].shape))
print("pts                 ", tuple(pts.shape))
print("ray 0, last point   ", pts[0, -1].tolist())

The same rules catch you on the way back. torch.sum(weights[..., None] * rgb, -2) is (R, N, 1) × (R, N, 3) summed over the samples: one colour per ray. Drop the [..., None] and (R, N) × (R, N, 3) aligns N with 3 — an error, unless N happens to be 3. Drop keepdim=True from a per-ray sum you divide by and (R, N) / (R,) aligns the rays with the samples: an error when R ≠ N, and silently wrong when R = N. Test code with square batches for exactly this reason.

2. Where to sample: stratified depths

nerf-pytorch builds t = linspace(0, 1, N) and z = near·(1 − t) + far·t: N depths from near to far including both ends. At test time (perturb = 0) those are the samples. During training each sample is jittered inside its stratum, bounded by the midpoints between neighbours, with fresh noise every iteration:

import torch

def stratified_samples(near, far, n, u=None):
    t = torch.linspace(0.0, 1.0, n, dtype=near.dtype)
    z = near[..., None] * (1.0 - t) + far[..., None] * t
    if u is None:
        return z
    mids = 0.5 * (z[..., 1:] + z[..., :-1])
    upper = torch.cat([mids, z[..., -1:]], -1)
    lower = torch.cat([z[..., :1], mids], -1)
    return lower + (upper - lower) * u

near, far = torch.tensor([2.0]), torch.tensor([6.0])
print("edges   ", stratified_samples(near, far, 5).tolist())
print("u = 0   ", stratified_samples(near, far, 5, torch.zeros(1, 5)).tolist())
print("u = 1   ", stratified_samples(near, far, 5, torch.ones(1, 5)).tolist())
torch.manual_seed(0)
print("jittered", [round(v, 3) for v in stratified_samples(near, far, 5, torch.rand(1, 5))[0].tolist()])

Why jitter? A fixed set of N depths lets the MLP be right only at those depths: it never sees anything in between, and novel views (whose samples land elsewhere) show banding. Jittering makes every training step evaluate a different quadrature of the same integral, so the network has to be right everywhere. Note the strata: the first and last are half as wide as the others, because the edges include near and far. The C++ skill's stratified_samples used N equal bins instead — both are fine; know which one the code you're porting uses.

3. raw2outputs: the pixel from the samples

The MLP returns raw, shape (R, N, 4): three colour logits and a raw density. Activations come first — sigmoid for colours in (0, 1), relu for σ ≥ 0 (mip-NeRF and later use softplus, smoother near zero; TensoRF and Instant-NGP use exp / truncated exp). Then the quadrature from the volume rendering integral (NeRF eq. 3):

• δ_i = z_{i+1} − z_i, with δ_{N−1} = 10¹⁰ (the last interval reaches to infinity), times ‖d‖;
• α_i = 1 − exp(−σ_i δ_i);
• T_i = Π_{j<i} (1 − α_j), the light that reaches sample i;
• w_i = T_i α_i; colour Σ w_i c_i, depth Σ w_i z_i, opacity ("acc") Σ w_i.

The ‖d‖ is easy to forget. nerf-pytorch's get_rays returns directions (x/f, −y/f, −1) — not unit length — and z_vals are parameters of o + t·d, so the world distance between samples is δ·‖d‖. Forget it and a corner pixel (longer d) sees a thinner medium than the centre pixel for the same field. nerf_pl normalises its directions instead, with near/far in world units: consistent too. Mixing the two conventions is the bug.

The exclusive product in one line

T_i is a running product that must exclude sample i itself. The C++ version kept a running T and updated it after computing the weight:

float T = 1.0f;
for (std::size_t i = 0; i < n; ++i) {
    w[i] = T * alpha[i];
    T *= 1.0f - alpha[i];   // after: T is exclusive
}

In PyTorch the loop is a scan, torch.cumprod — but cumprod is inclusive: its i-th output contains factor i. Shift it one place right by prepending a column of ones and dropping the last column:

import math
import torch

def raw2outputs(raw, z_vals, rays_d, white_bkgd=False):
    rgb = torch.sigmoid(raw[..., :3])
    sigma = torch.relu(raw[..., 3])
    dists = z_vals[..., 1:] - z_vals[..., :-1]
    dists = torch.cat([dists, torch.full_like(z_vals[..., :1], 1e10)], -1)
    dists = dists * torch.linalg.vector_norm(rays_d, dim=-1, keepdim=True)
    alpha = 1.0 - torch.exp(-sigma * dists)
    ones = torch.ones_like(alpha[..., :1])
    T = torch.cumprod(torch.cat([ones, 1.0 - alpha + 1e-10], -1), -1)[..., :-1]
    weights = alpha * T
    rgb_map = torch.sum(weights[..., None] * rgb, -2)
    depth_map = torch.sum(weights * z_vals, -1)
    acc_map = torch.sum(weights, -1)
    if white_bkgd:
        rgb_map = rgb_map + (1.0 - acc_map[..., None])
    return rgb_map, depth_map, acc_map, weights

# one ray, sigma*delta = ln 2 on the first three samples: alpha = 0.5 each
s = math.log(2.0)
raw = torch.tensor([[[3.0, 0.0, -3.0, s], [0.0, 3.0, 0.0, s], [0.0, 0.0, 3.0, s], [0.0, 0.0, 0.0, -1.0]]])
z = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
rgb, depth, acc, w = raw2outputs(raw, z, torch.tensor([[0.0, 0.0, 1.0]]))
print("weights", [round(v, 4) for v in w[0].tolist()])
print("acc", round(acc.item(), 4), "depth", round(depth.item(), 4), "rgb", [round(v, 3) for v in rgb[0].tolist()])
rgb_w, *_ = raw2outputs(raw, z, torch.tensor([[0.0, 0.0, 1.0]]), white_bkgd=True)
print("on white", [round(v, 3) for v in rgb_w[0].tolist()])
# the same ray with an unnormalised direction: twice the distance, more opacity
_, _, acc2, _ = raw2outputs(raw, z, torch.tensor([[0.0, 0.0, 2.0]]))
print("|d| = 2 -> acc", round(acc2.item(), 4))

Three details from that code:

1. + 1e-10 inside the product. Once a sample is fully opaque, 1 − α is exactly 0, every later T is exactly 0, and so is every gradient to samples behind it. The epsilon (carried over from the TF original) keeps the product from collapsing to an exact zero. nerfstudio avoids the issue by computing T in log space: T_i = exp(−Σ_{j<i} σ_j δ_j) with a cumsum, which is the same number (1 − α_j = e^{−σ_j δ_j}) and better conditioned.
2. depth is not divided by acc. On a ray that hits nothing, depth is ≈ 0. nerf-pytorch's disp_map = 1 / max(1e-10, depth / acc) does the division, guarded.
3. White background: NeRF-synthetic images are RGBA, composited onto white for training. The renderer must then also composite onto white: the light that wasn't absorbed, 1 − acc, is added to every channel of that ray — (1 - acc_map)[..., None], shape (R, 1) against (R, 3).

4. Differentiable means: no in-place, no stray no_grad

Training backpropagates the photometric loss through raw2outputs into the MLP, so every op above must be recorded by autograd. Two rules keep that true. First, don't write into tensors autograd may need: torch.cumprod saves its output for the backward pass, so "shifting it" with T[..., 1:] = T[..., :-1].clone(); T[..., 0] = 1 computes the right values and then fails in backward():

import torch

alpha = torch.tensor([0.3, 0.6, 0.9], requires_grad=True)
T = torch.cumprod(1.0 - alpha, -1)
T[1:] = T[:-1].clone()    # looks harmless
T[0] = 1.0
try:
    (alpha * T).sum().backward()
except RuntimeError as e:
    print("backward:", str(e).split(":")[0])

alpha.grad = None
T = torch.cumprod(torch.cat([torch.ones(1), 1.0 - alpha]), -1)[:-1]   # out of place
(alpha * T).sum().backward()
print("grad", [round(v, 3) for v in alpha.grad.tolist()])

Every tensor has a version counter; in-place ops bump it, and backward checks that what it saved is unchanged. Second, torch.no_grad() and .detach() belong exactly where you mean "no gradient here". In NeRF that is one place: the positions of the fine samples (section 5). The coarse pass is trained by its own loss term, so wrapping it in no_grad to "save memory" silently freezes the coarse network.

Check a differentiable function the way PyTorch's own tests do, with torch.autograd.gradcheck in float64 — it compares your backward pass with finite differences:

import torch

def composite_weights(sigma, z):
    dists = torch.cat([z[..., 1:] - z[..., :-1], torch.full_like(z[..., :1], 1e10)], -1)
    alpha = 1.0 - torch.exp(-torch.relu(sigma) * dists)
    T = torch.cumprod(torch.cat([torch.ones_like(alpha[..., :1]), 1.0 - alpha + 1e-10], -1), -1)[..., :-1]
    return alpha * T

torch.manual_seed(0)
sigma = (torch.rand(3, 5, dtype=torch.float64) + 0.1).requires_grad_()
z = torch.sort(torch.rand(3, 5, dtype=torch.float64) * 4, -1).values
print("gradcheck:", torch.autograd.gradcheck(lambda s: composite_weights(s, z), (sigma,)))

In float32 the finite differences are too noisy for gradcheck's tolerances — that is why it is always run in float64, and why your functions must keep whatever dtype they are given (full_like, ones_like, dtype= on linspace), rather than quietly producing float32.

5. Hierarchical sampling: sample_pdf

The coarse pass wastes most samples on empty space. NeRF reuses its weights as a histogram over depth: normalise them into a piecewise-constant pdf over the bins between coarse samples, and draw N_f more samples from it by inverse transform sampling — take u ∈ [0, 1], find where the CDF crosses u, interpolate inside that bin.

The search is torch.searchsorted, batched: for each ray, the indices where its u values would be inserted into its sorted CDF. With right=True it returns the first index whose CDF value is greater than u, so inds − 1 is the bin u falls in:

import torch

def sample_pdf(bins, weights, n_samples, u=None):
    weights = weights.detach() + 1e-5                     # positions, not a gradient path
    pdf = weights / torch.sum(weights, -1, keepdim=True)
    cdf = torch.cat([torch.zeros_like(pdf[..., :1]), torch.cumsum(pdf, -1)], -1)
    if u is None:
        u = torch.linspace(0.0, 1.0, n_samples, dtype=cdf.dtype).expand(cdf.shape[:-1] + (n_samples,))
    u = u.contiguous()
    inds = torch.searchsorted(cdf, u, right=True)
    below = torch.clamp(inds - 1, min=0)
    above = torch.clamp(inds, max=cdf.shape[-1] - 1)
    cdf_b, cdf_a = torch.gather(cdf, -1, below), torch.gather(cdf, -1, above)
    bins_b, bins_a = torch.gather(bins, -1, below), torch.gather(bins, -1, above)
    denom = torch.where(cdf_a - cdf_b < 1e-5, torch.ones_like(cdf_a), cdf_a - cdf_b)
    return bins_b + (u - cdf_b) / denom * (bins_a - bins_b)

bins = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0]])
w = torch.tensor([[0.05, 0.1, 0.7, 0.15]], requires_grad=True)
s = sample_pdf(bins, w, 9, torch.linspace(0.1, 0.9, 9)[None])
print([round(v, 2) for v in s[0].tolist()], "requires_grad:", s.requires_grad)
print("all-zero weights ->", [round(v, 2) for v in sample_pdf(bins, torch.zeros(1, 4), 5)[0].tolist()])

The guards each matter. + 1e-5 turns all-zero weights (a ray through empty space) into a uniform pdf instead of 0/0. The clamps keep gather in range at u = 0 and u = 1. A bin of (almost) zero probability would divide by ≈ 0, so tiny denominators become 1. contiguous() because expand produces a stride-0 view and searchsorted warns (and copies) on non-contiguous input.

In render_rays, the bins are the midpoints of the coarse z values and the weights are weights[..., 1:-1] — the first and last coarse samples have no bin of their own. The new samples are concatenated with the coarse ones and sorted (torch.sort(torch.cat([z_vals, z_samples], -1), -1)), and the fine network evaluates all N_c + N_f points. The positions are detached: gradients reach the field through what was seen at those depths, never through the choice of depths.

6. Positional encoding

A ReLU MLP fed raw (x, y, z) learns only smooth functions — the spectral bias that Tancik et al. analyse in Fourier Features Let Networks Learn High Frequency Functions. NeRF lifts each coordinate to sines and cosines of octave frequencies first: γ(p) = (sin(2⁰πp), cos(2⁰πp), …, sin(2^{L−1}πp), cos(2^{L−1}πp)), L = 10 for positions and L = 4 for directions.

import torch

def positional_encoding(x, n_freqs, include_input=True):
    out = [x] if include_input else []
    for k in range(n_freqs):
        out += [torch.sin(x * 2.0**k), torch.cos(x * 2.0**k)]
    return torch.cat(out, -1)

x = torch.tensor([[0.5, -1.0, 2.0]])
e = positional_encoding(x, 2)
print(tuple(e.shape), [round(v, 3) for v in e[0].tolist()])
print("pts:", positional_encoding(torch.zeros(1024, 64, 3), 10).shape[-1],
      "dirs:", positional_encoding(torch.zeros(1024, 64, 3), 4).shape[-1])

That's 63 and 27 input features — the numbers hard-coded in the NeRF MLP. And here the codebases disagree in ways that only matter when you load someone else's weights:

The missing π in nerf-pytorch is harmless (the scene is scaled to roughly [−1, 1]³ or [−4, 4]³ and the frequencies are just a factor π lower); the layout difference silently scrambles a ported checkpoint. Build the frequencies as Python floats (2.0**k), not as a NumPy array: a float64 tensor from torch.from_numpy would promote a float32 encoding to float64.

7. Chunks: bounding peak memory

One training batch is R = 1024 rays × 192 samples ≈ 200 k points through eight 256-wide layers — fine. A test image is 640 000 rays: ≈ 123 M points, and one layer's activations alone would be ≈ 126 GB. nerf-pytorch therefore evaluates the network with batchify(fn, chunk) and renders with batchify_rays(rays, chunk): slice along dim 0, run, concatenate.

import torch

def run_chunked(fn, x, chunk):
    if x.shape[0] == 0:
        return fn(x)               # torch.cat([]) would raise
    return torch.cat([fn(x[i:i + chunk]) for i in range(0, x.shape[0], chunk)], 0)

torch.manual_seed(0)
mlp = torch.nn.Sequential(torch.nn.Linear(63, 64), torch.nn.ReLU(), torch.nn.Linear(64, 4))
x = torch.randn(10_000, 63)
full, chunked = mlp(x), run_chunked(mlp, x, 4096)
print("same output:", torch.allclose(full, chunked, atol=1e-6))
g1 = torch.autograd.grad(full.square().mean(), mlp[0].weight)[0]
g2 = torch.autograd.grad(chunked.square().mean(), mlp[0].weight)[0]
print("same gradient:", torch.allclose(g1, g2, atol=1e-6))

Chunking changes neither the result nor the gradients — every point is processed independently and torch.cat is differentiable — only how much lives in memory at once. At test time, add torch.no_grad() as well, so the activations of each chunk are freed as soon as the next layer has run.

8. The dtype and shape checklist

The bugs that reach a NeRF code review are rarely in the formulas. They are:

• a reduction over the wrong axis, or a missing keepdim=True — test with a square batch (as many rays as samples);
• an inclusive cumprod (the sample hides itself) or an in-place shift of it;
• δ without ‖d‖, or directions normalised while near/far are not;
• float64 sneaking in from NumPy (torch.from_numpy, torch.tensor(np_array)) or float32 from factories (torch.linspace, torch.ones without dtype=);
• no_grad around the coarse pass, or a missing detach on the fine samples;
• torch.cat of an empty list for an empty batch.

And the conventions that differ between codebases, which you should check before porting anything:

The build task asks for all of it, nerf-pytorch flavoured: stratified samples, raw2outputs, sample_pdf, the encoding, chunking, and a two-pass render_rays — vectorised over any batch shape and checked with gradcheck.