别再死磕公式了！用PyTorch从零实现一个NeRF，带你直观理解神经辐射场

从零构建NeRF：用PyTorch实现神经辐射场的实战指南

如果你已经厌倦了那些充满复杂数学推导的NeRF教程，那么这篇文章正是为你准备的。我们将完全从代码角度出发，用PyTorch一步步实现一个精简但功能完整的NeRF模型。不需要深厚的数学背景，只要你对3D视觉和深度学习有基本了解，就能跟随这个教程构建出自己的神经辐射场渲染器。

1. 环境准备与基础架构

在开始之前，确保你的开发环境已经安装了PyTorch 1.8+和相关的视觉处理库。我们将从构建最基础的MLP网络开始，这是NeRF的核心组件。

import torch import torch.nn as nn import torch.nn.functional as F class BasicNeRF(nn.Module): def __init__(self, pos_enc_dim=10, dir_enc_dim=4, hidden_dim=128): super().__init__() # 位置编码后的输入维度 pos_enc_in_dim = 3 * (2 * pos_enc_dim + 1) dir_enc_in_dim = 3 * (2 * dir_enc_dim + 1) # 处理位置的主网络 self.block1 = nn.Linear(pos_enc_in_dim, hidden_dim) self.block2 = nn.Linear(hidden_dim, hidden_dim) self.block3 = nn.Linear(hidden_dim, hidden_dim) # 密度输出层 self.density_out = nn.Linear(hidden_dim, 1) # 处理方向的第二网络 self.block4 = nn.Linear(hidden_dim + dir_enc_in_dim, hidden_dim//2) self.color_out = nn.Linear(hidden_dim//2, 3) def forward(self, pos, dir): # 实现将在下一节填充 pass

这个基础架构包含了NeRF的两个关键部分：

• 位置处理网络：将3D坐标转换为体积密度和中间特征
• 方向处理网络：结合视角方向预测RGB颜色

注意：这里的pos_enc_dim和dir_enc_dim控制位置编码的维度，我们将在下一节详细讨论这个关键组件。

2. 实现位置编码：打破低频瓶颈

位置编码(PE)是NeRF能够捕捉高频细节的关键。传统的MLP网络倾向于学习低频函数，而PE通过将输入映射到高维空间，使网络更容易学习高频变化。

def positional_encoding(x, L): """位置编码函数 Args: x: 输入坐标/方向 (..., 3) L: 编码级别 Returns: encoded: 编码后的向量 (..., 3*(2L+1)) """ encodings = [x] for i in range(L): for fn in [torch.sin, torch.cos]: encodings.append(fn(2.**i * x)) return torch.cat(encodings, dim=-1)

这个编码函数会产生以下形式的输出：

• 原始坐标 (x,y,z)
• sin(2⁰πx), cos(2⁰πx), sin(2⁰πy), cos(2⁰πy), sin(2⁰πz), cos(2⁰πz)
• sin(2¹πx), cos(2¹πx), ..., 直到2ᴸ⁻¹π

现在我们可以完善BasicNeRF的forward方法：

def forward(self, pos, dir): # 位置编码 pos_encoded = positional_encoding(pos, self.pos_enc_dim) dir_encoded = positional_encoding(dir, self.dir_enc_dim) # 通过主网络处理位置 h = F.relu(self.block1(pos_encoded)) h = F.relu(self.block2(h)) h = F.relu(self.block3(h)) # 预测密度 density = F.relu(self.density_out(h)) # 结合方向预测颜色 h = torch.cat([h, dir_encoded], dim=-1) h = F.relu(self.block4(h)) color = torch.sigmoid(self.color_out(h)) # sigmoid确保颜色在[0,1] return color, density

3. 体渲染实现：从神经网络输出到图像

有了NeRF网络，我们需要实现体渲染算法将预测的密度和颜色转换为像素值。这是NeRF最复杂的部分之一，但我们可以分步实现。

首先，定义光线生成函数：

def generate_rays(height, width, focal, pose): """生成相机光线 Args: height, width: 图像尺寸 focal: 相机焦距 pose: 相机位姿 (4x4矩阵) Returns: rays_o: 光线原点 (H, W, 3) rays_d: 光线方向 (H, W, 3) """ i, j = torch.meshgrid(torch.arange(width), torch.arange(height)) directions = torch.stack([(i-width/2)/focal, -(j-height/2)/focal, -torch.ones_like(i)], -1) rays_d = torch.sum(directions[..., None, :] * pose[:3, :3], -1) rays_o = pose[:3, -1].expand(rays_d.shape) return rays_o, rays_d

接下来是核心的体渲染函数：

def volume_render(rays_o, rays_d, near, far, model, num_samples=64): """体渲染函数 Args: rays_o: 光线原点 (N_rays, 3) rays_d: 光线方向 (N_rays, 3) near/far: 采样范围 model: NeRF模型 num_samples: 每条光线的采样点数 Returns: rgb_map: 渲染的RGB图像 (N_rays, 3) depth_map: 深度图 (N_rays,) """ # 1. 计算采样点 t_vals = torch.linspace(near, far, num_samples) points = rays_o[..., None, :] + rays_d[..., None, :] * t_vals[..., None] # 2. 扩展方向向量以匹配采样点 dirs = rays_d[..., None, :].expand(points.shape) # 3. 展平以批量处理 points_flat = points.reshape(-1, 3) dirs_flat = dirs.reshape(-1, 3) # 4. 通过网络获取颜色和密度 rgb_flat, density_flat = model(points_flat, dirs_flat) rgb = rgb_flat.reshape(list(points.shape[:-1]) + [3]) density = density_flat.reshape(list(points.shape[:-1])) # 5. 计算透明度 delta = t_vals[..., 1:] - t_vals[..., :-1] delta = torch.cat([delta, torch.tensor([1e10]).expand(delta[..., :1].shape)], -1) alpha = 1 - torch.exp(-density * delta) # 6. 计算权重 (transmittance * alpha) trans = torch.exp(-torch.cat([torch.zeros_like(density[..., :1]), torch.cumsum(density[..., :-1] * delta[..., :-1], -1)], -1)) weights = trans * alpha # 7. 合成最终颜色和深度 rgb_map = torch.sum(weights[..., None] * rgb, -2) depth_map = torch.sum(weights * t_vals, -1) return rgb_map, depth_map

这个渲染过程实现了NeRF论文中的关键公式，但完全用PyTorch操作表达，避免了复杂的数学符号。

4. 训练策略与技巧

训练NeRF模型需要一些特别的技巧才能获得好结果。以下是经过实践验证的有效方法：

4.1 分层采样

原始NeRF论文提出了分层采样策略，先在光线粗采样，然后根据密度分布进行精细采样。我们可以这样实现：

def hierarchical_sampling(rays_o, rays_d, near, far, model, num_coarse=64, num_fine=128): # 粗采样 t_vals_coarse = torch.linspace(near, far, num_coarse) points_coarse = rays_o[..., None, :] + rays_d[..., None, :] * t_vals_coarse[..., None] # 获取粗采样密度 with torch.no_grad(): _, density_coarse = model(points_coarse, rays_d[..., None, :].expand(points_coarse.shape)) # 根据密度分布生成精细采样点 weights = density_coarse * (t_vals_coarse[..., 1:] - t_vals_coarse[..., :-1]) pdf = weights / (torch.sum(weights, -1, keepdim=True) + 1e-5) cdf = torch.cumsum(pdf, -1) cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1) # 逆变换采样 u = torch.rand(list(cdf.shape[:-1]) + [num_fine]) idx = torch.searchsorted(cdf, u, right=True) lower = torch.max(torch.zeros_like(idx), idx-1) upper = torch.min(torch.ones_like(idx)*(num_coarse-1), idx) idx_g = torch.stack([lower, upper], -1) cdf_g = torch.gather(cdf, -1, idx_g) t_vals_coarse_g = torch.gather(t_vals_coarse, -1, idx_g) denom = cdf_g[..., 1] - cdf_g[..., 0] denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom) t = (u - cdf_g[..., 0]) / denom t_fine = t_vals_coarse_g[..., 0] + t * (t_vals_coarse_g[..., 1] - t_vals_coarse_g[..., 0]) # 合并粗采样和精细采样点 t_vals, _ = torch.sort(torch.cat([t_vals_coarse, t_fine], -1), -1) points = rays_o[..., None, :] + rays_d[..., None, :] * t_vals[..., None] return points, t_vals

4.2 训练循环实现

完整的训练循环需要考虑以下几个关键点：

def train(): # 初始化模型和优化器 model = BasicNeRF().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=5e-4) # 加载数据集 (这里假设已有数据加载器) dataset = load_dataset() dataloader = DataLoader(dataset, batch_size=1024, shuffle=True) for epoch in range(1000): for batch in dataloader: rays_o, rays_d, target_rgb = batch # 前向传播 rgb_coarse, _ = volume_render(rays_o, rays_d, 2., 6., model, 64) # 分层采样 points_fine, t_vals_fine = hierarchical_sampling(rays_o, rays_d, 2., 6., model) dirs_fine = rays_d[..., None, :].expand(points_fine.shape) rgb_fine, _ = volume_render(rays_o, rays_d, 2., 6., model, 128) # 计算损失 loss_coarse = F.mse_loss(rgb_coarse, target_rgb) loss_fine = F.mse_loss(rgb_fine, target_rgb) loss = loss_coarse + loss_fine # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() # 每隔一定epoch保存模型和渲染测试视图 if epoch % 50 == 0: render_test_view(model) torch.save(model.state_dict(), f"model_{epoch}.pth")

4.3 实用技巧与调优

在实际训练中，以下几个技巧能显著提升模型性能：

1. 学习率调度：

   scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)

2. 权重初始化：

   def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) nn.init.zeros_(m.bias) model.apply(init_weights)

3. 颜色空间转换：

   • 在输入网络前对颜色进行gamma校正
   • 使用线性RGB空间进行训练

4. 正则化技巧：

   • 对密度预测添加L2正则
   • 使用权重衰减防止过拟合

5. 可视化与结果分析

训练完成后，我们可以渲染不同视角的图像来评估模型性能。以下是几个关键的可视化方法：

5.1 新视角合成

def render_novel_views(model, num_views=30): for i in range(num_views): # 生成环绕相机轨迹 theta = 2 * np.pi * i / num_views pose = create_camera_pose(theta) # 生成光线 rays_o, rays_d = generate_rays(400, 400, 500, pose) # 渲染图像 with torch.no_grad(): rgb, _ = volume_render(rays_o, rays_d, 2., 6., model, 128) # 保存或显示结果 save_image(rgb, f"view_{i}.png")

5.2 深度图可视化

深度图可以帮助我们理解场景的几何结构：

def visualize_depth(depth_map): depth_map = depth_map - depth_map.min() depth_map = depth_map / depth_map.max() depth_map = 1 - depth_map # 近处亮，远处暗 plt.imshow(depth_map, cmap='viridis') plt.colorbar() plt.show()

5.3 性能评估指标

定量评估可以使用以下指标：

指标名称	计算公式	说明
PSNR	10·log₁₀(MAX²/MSE)	峰值信噪比，值越高越好
SSIM	结构相似性指数	衡量结构相似性，范围[0,1]
LPIPS	感知相似性	基于深度学习的感知指标

实现示例：

def compute_metrics(pred, target): mse = F.mse_loss(pred, target) psnr = -10. * torch.log10(mse) # SSIM计算 (需要实现或使用现有库) ssim = compute_ssim(pred, target) return {"PSNR": psnr.item(), "SSIM": ssim}

6. 扩展与优化方向

基础NeRF实现完成后，可以考虑以下几个优化方向：

6.1 加速渲染

原始NeRF渲染速度很慢，可以考虑：

1. 空间数据结构：

   • 使用八叉树或KD树加速空间查询
   • 实现重要性采样减少无效计算

2. 网络架构优化：

   class EfficientNeRF(nn.Module): def __init__(self): super().__init__() # 实现更高效的网络结构 pass

6.2 处理动态场景

扩展NeRF处理动态场景：

class DynamicNeRF(nn.Module): def __init__(self): super().__init__() # 添加时间维度处理 self.time_encoder = nn.Linear(1, 64) def forward(self, pos, dir, time): # 结合时间信息 time_enc = self.time_encoder(time) # 其余处理类似基础NeRF

6.3 真实场景应用

在实际应用中还需要考虑：

1. 数据预处理：

   • 相机标定与位姿估计
   • 光照一致性处理

2. 大规模场景：

   • 分块训练策略
   • 多分辨率表示

3. 实时渲染：

   • 网络蒸馏
   • 预计算技术

7. 常见问题与解决方案

在实现NeRF过程中，你可能会遇到以下典型问题：

问题现象	可能原因	解决方案
渲染结果全黑	密度预测过大	降低密度预测的初始偏置
颜色过饱和	输出激活函数不当	使用sigmoid限制颜色范围
细节丢失	位置编码不足	增加位置编码维度L
训练不稳定	学习率过高	使用学习率调度器
渲染伪影	采样点不足	增加采样点或使用分层采样

对于更复杂的问题，可以尝试：

# 梯度裁剪防止爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 混合精度训练 scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): rgb, _ = volume_render(rays_o, rays_d, 2., 6., model)

8. 完整代码结构与项目组织

一个结构良好的NeRF项目可以这样组织：

nerf-project/ ├── configs/ # 配置文件 │ └── base.yaml ├── data/ # 数据集 │ └── synthetic/ ├── models/ # 模型定义 │ ├── __init__.py │ ├── basic_nerf.py # 基础模型 │ └── efficient.py # 优化模型 ├── render/ # 渲染相关 │ ├── rays.py # 光线生成 │ └── volume.py # 体渲染 ├── utils/ # 工具函数 │ ├── metrics.py # 评估指标 │ └── visualization.py # 可视化 ├── train.py # 训练脚本 └── render.py # 渲染脚本

关键训练脚本结构：

# train.py def main(config): # 初始化 model = build_model(config) optimizer = build_optimizer(model, config) dataloader = build_dataloader(config) # 训练循环 for epoch in range(config.epochs): for batch in dataloader: # 训练步骤 train_step(model, optimizer, batch) # 验证和保存 if epoch % config.val_freq == 0: validate(model, val_loader) save_checkpoint(model, epoch)

通过这样的项目结构，你可以轻松扩展和维护代码，尝试不同的NeRF变体和改进方案。