从NeRF到NeuS：手把手教你用PyTorch复现SDF体渲染，搞定多视角三维重建

从NeRF到NeuS：用PyTorch实现SDF体渲染的工程实践

三维重建技术正在经历从传统多视图几何到神经隐式表达的范式转移。2020年NeRF的横空出世，展示了神经辐射场在视角合成方面的惊人能力，但其密度场表示难以直接提取高质量表面。2021年NeuS的提出，通过将符号距离函数（SDF）与体渲染相结合，实现了无需mask监督的高保真表面重建。本文将深入解析NeuS的核心算法，并手把手教你用PyTorch实现这套创新性的SDF体渲染系统。

1. 环境准备与依赖安装

1.1 基础环境配置

推荐使用Python 3.8+和PyTorch 1.10+环境，以下是关键依赖的安装命令：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 pip install numpy scipy matplotlib tqdm opencv-python imageio

对于GPU加速，需要确保CUDA工具包版本与PyTorch匹配。可以通过以下代码验证环境：

import torch print(f"PyTorch版本: {torch.__version__}") print(f"CUDA可用: {torch.cuda.is_available()}") print(f"GPU型号: {torch.cuda.get_device_name(0)}")

1.2 项目结构设计

建议采用如下模块化结构组织代码：

neus-pytorch/ ├── core/ # 核心算法实现 │ ├── networks.py # 神经网络定义 │ ├── rendering.py # 渲染方程实现 │ └── sdf_utils.py # SDF相关工具 ├── data/ # 数据加载与处理 ├── configs/ # 训练配置文件 ├── outputs/ # 训练结果输出 └── train.py # 主训练脚本

2. NeuS核心算法解析

2.1 S-density的数学本质

NeuS的核心创新在于提出了S-density函数，将SDF值与体积渲染概率密度建立联系：

def phi_s(x, s): """S-density函数实现""" return s * torch.exp(-s*x) / (1 + torch.exp(-s*x))**2

这个函数实际上是sigmoid函数的导数，具有以下重要性质：

• 当SDF值f(x)→0（接近表面）时，ϕₛ(f(x))取得最大值
• 参数s控制分布尖锐程度，s越大分布越集中
• 在训练过程中，s会自适应调整以提高重建精度

2.2 无偏权重函数设计

NeuS设计了满足两个关键性质的权重函数：

1. 无偏性：在表面交点处权重最大
2. 遮挡感知：近处表面贡献大于远处

实现代码的核心部分：

def compute_weights(sdf, z_vals, cos_theta, s): """ 计算满足无偏和遮挡感知的渲染权重 参数: sdf: [N_rays, N_samples] 采样点的SDF值 z_vals: [N_rays, N_samples] 采样深度 cos_theta: [N_rays, 1] 视线与表面法向夹角余弦 s: 可训练的比例参数 返回: weights: [N_rays, N_samples] 渲染权重 """ # 计算S-density density = phi_s(sdf, s) # 计算不透明密度ρ(t) cdf = torch.sigmoid(sdf * s) delta = z_vals[:, 1:] - z_vals[:, :-1] delta = torch.cat([delta, delta[:, -1:]], dim=-1) rho = (cdf[:, :-1] - cdf[:, 1:]) / (cdf[:, :-1] + 1e-5) rho = rho.clamp(min=0) / delta[:, :-1] # 计算累积透射率 T = torch.exp(-torch.cumsum(rho * delta[:, :-1], dim=-1)) T = torch.cat([torch.ones_like(T[:, :1]), T], dim=-1) # 最终权重 weights = T * density * cos_theta.abs() return weights / (weights.sum(-1, keepdim=True) + 1e-5)

3. 网络架构实现

3.1 SDF网络设计

NeuS采用8层MLP预测SDF值，关键实现细节：

class SDFNetwork(nn.Module): def __init__(self): super().__init__() self.skip_in = [4] # 在第4层添加跳跃连接 self.softplus = nn.Softplus(beta=100) # 位置编码 self.embed_fn, input_ch = get_embedder(6) layers = [] dims = [input_ch] + [256] * 8 for i in range(8): if i in self.skip_in: dims[i+1] += input_ch layers.append(nn.Linear(dims[i], dims[i+1])) if i != 7: # 最后一层不加激活 layers.append(self.softplus) self.layers = nn.ModuleList(layers) def forward(self, x): x_embed = self.embed_fn(x) h = x_embed for i, l in enumerate(self.layers): if i in self.skip_in: h = torch.cat([h, x_embed], -1) h = l(h) return h

3.2 颜色网络设计

颜色网络考虑视角相关效果，实现如下：

class ColorNetwork(nn.Module): def __init__(self): super().__init__() self.embed_fn, input_ch = get_embedder(6) # 位置编码 self.embed_dir_fn, dir_ch = get_embedder(4) # 方向编码 # 网络结构 self.fc1 = nn.Linear(input_ch + 256 + dir_ch + 3, 256) self.fc2 = nn.Linear(256, 256) self.fc3 = nn.Linear(256, 3) self.act = nn.ReLU() def forward(self, x, d, sdf_feat, normal): x_embed = self.embed_fn(x) d_embed = self.embed_dir_fn(d) h = torch.cat([x_embed, sdf_feat, d_embed, normal], -1) h = self.act(self.fc1(h)) h = self.act(self.fc2(h)) return torch.sigmoid(self.fc3(h)) # 输出RGB颜色

4. 训练流程与优化技巧

4.1 分层采样策略

NeuS采用类似NeRF的粗-细采样策略，但有以下改进：

1. 粗采样阶段使用固定s值计算权重
2. 细采样阶段使用网络预测的s值
3. 在表面附近集中采样点

实现代码框架：

def sample_pdf(bins, weights, N_samples): """基于权重分布的重要性采样""" weights = weights + 1e-5 pdf = weights / weights.sum(-1, keepdim=True) cdf = torch.cumsum(pdf, -1) # 生成均匀采样点 u = torch.linspace(0, 1, N_samples) u = u.expand(list(cdf.shape[:-1]) + [N_samples]) # 反变换采样 idx = torch.searchsorted(cdf, u, right=True) below = torch.max(torch.zeros_like(idx-1), idx-1) above = torch.min((cdf.shape[-1]-1)*torch.ones_like(idx), idx) idx_g = torch.stack([below, above], -1) # 线性插值 cdf_g = torch.gather(cdf, -1, idx_g) bins_g = torch.gather(bins, -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 samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0]) return samples

4.2 损失函数设计

NeuS的损失函数包含三个关键部分：

def compute_loss(rgb_pred, rgb_gt, sdf, pts, mask_pred=None, mask_gt=None): # 颜色损失 color_loss = F.mse_loss(rgb_pred, rgb_gt) # SDF正则化 (Eikonal项) grad = torch.autograd.grad(sdf.sum(), pts, create_graph=True)[0] eikonal_loss = ((grad.norm(2, dim=-1) - 1)**2).mean() # 可选mask损失 mask_loss = 0 if mask_pred is not None and mask_gt is not None: mask_loss = F.binary_cross_entropy(mask_pred, mask_gt) return { 'color': color_loss, 'eikonal': eikonal_loss, 'mask': mask_loss, 'total': color_loss + 0.1*eikonal_loss + mask_loss }

5. 实战调试与性能优化

5.1 训练技巧

1. 学习率调度：采用余弦退火策略

   scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=300000, eta_min=1e-5)

2. 梯度裁剪：防止SDF网络梯度爆炸

   torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

3. 自适应s参数：初始值设为0.3，随训练自动调整

5.2 常见问题排查

5.3 性能优化策略

1. 混合精度训练：

   scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): # 前向计算... scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

2. 射线批处理：将射线分组处理以减少内存占用

3. 提前体素剔除：使用八叉树加速空区域跳过

6. 结果可视化与分析

6.1 定量评估指标

在DTU数据集上的典型评估结果：

CD为倒角距离(×10³)，F-score为F1分数(%)

6.2 定性结果对比

NeuS相比NeRF在以下方面表现更优：

• 几何细节：能重建更尖锐的边缘和薄结构
• 表面连续性：减少漂浮物和空洞现象
• 纹理保真度：保持高频纹理细节

6.3 局限性与改进方向

当前实现仍存在以下挑战：

1. 计算成本：每场景仍需数小时训练
2. 纹理复制：在无纹理区域可能出现伪影
3. 尺度敏感：对场景尺度假设较强

可能的改进方向包括：

• 采用哈希编码加速训练
• 引入显式几何先验
• 开发多尺度表示方法