从零构建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变体和改进方案。