用PyTorch复现NeRF：从零开始手把手教你训练自己的乐高小车3D模型（附完整代码）

用PyTorch复现NeRF从零开始手把手教你训练自己的乐高小车3D模型附完整代码当你第一次看到NeRF生成的3D模型时那种震撼感是难以言喻的——从一个简单的神经网络中竟然能重建出如此精细的三维场景。作为计算机视觉领域近年来的重大突破神经辐射场技术正在改变我们对3D重建的认知。本文将带你从零开始用PyTorch完整实现一个NeRF模型并以经典的乐高小车数据集为例一步步教你训练出自己的3D重建模型。1. 环境准备与数据加载在开始之前我们需要搭建一个合适的开发环境。建议使用Python 3.8和PyTorch 1.10版本这些组合经过验证能够很好地支持NeRF的实现。基础环境配置conda create -n nerf python3.8 conda activate nerf pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 pip install numpy matplotlib imageio scikit-image opencv-python乐高数据集是NeRF论文中使用的标准数据集之一它包含了从不同角度拍摄的乐高小车图像及对应的相机参数。我们可以从官方仓库下载这个数据集import os import json import numpy as np import imageio def load_lego_data(data_dir): 加载乐高数据集 with open(os.path.join(data_dir, transforms_train.json), r) as f: train_meta json.load(f) images [] poses [] for frame in train_meta[frames]: img_path os.path.join(data_dir, frame[file_path][2:]) images.append(imageio.imread(img_path)) poses.append(np.array(frame[transform_matrix])) images np.stack(images, axis0) poses np.stack(poses, axis0) return images, poses, train_meta[camera_angle_x]注意数据集中的图像需要归一化到0-1范围相机参数需要正确处理世界坐标系和相机坐标系之间的转换关系。2. NeRF模型架构实现NeRF的核心是一个多层感知机(MLP)它接收5D坐标(3D位置2D视角方向)作为输入输出体积密度和视角相关的RGB颜色值。让我们用PyTorch来实现这个网络import torch import torch.nn as nn import torch.nn.functional as F class PositionalEncoding(nn.Module): 位置编码模块将低维输入映射到高维空间 def __init__(self, L10): super().__init__() self.L L def forward(self, x): encodings [] for i in range(self.L): encodings.append(torch.sin(2**i * torch.pi * x)) encodings.append(torch.cos(2**i * torch.pi * x)) return torch.cat(encodings, dim-1) class NeRF(nn.Module): NeRF核心网络结构 def __init__(self, L_pos10, L_dir4): super().__init__() # 位置编码 self.pos_encoder PositionalEncoding(L_pos) self.dir_encoder PositionalEncoding(L_dir) # 主干网络 self.backbone nn.Sequential( nn.Linear(3 3*2*L_pos, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), ) # 密度预测头 self.density_head nn.Sequential( nn.Linear(256, 1), nn.ReLU() ) # 颜色预测头 self.color_head nn.Sequential( nn.Linear(256 3*2*L_dir, 128), nn.ReLU(), nn.Linear(128, 3), nn.Sigmoid() ) def forward(self, x, d): # 编码位置和方向 x_encoded self.pos_encoder(x) d_encoded self.dir_encoder(d) # 通过主干网络 features self.backbone(x_encoded) # 预测密度 sigma self.density_head(features) # 预测颜色 color_features torch.cat([features, d_encoded], dim-1) rgb self.color_head(color_features) return rgb, sigma提示位置编码是NeRF能够捕捉高频细节的关键确保L_pos和L_dir参数设置与论文一致。3. 体渲染实现与训练流程有了模型架构后我们需要实现NeRF的核心创新之一——可微分的体渲染过程。这个过程将模型预测的密度和颜色转换为最终的2D图像。def render_rays(model, rays_o, rays_d, near, far, N_samples, randFalse): 渲染单条射线上的颜色 # 计算采样点 t_vals torch.linspace(near, far, N_samples) if rand: t_vals t_vals torch.rand_like(t_vals) * (far-near)/N_samples # 采样点坐标 pts rays_o[...,None,:] rays_d[...,None,:] * t_vals[...,:,None] # 扩展视角方向 dirs rays_d / torch.norm(rays_d, dim-1, keepdimTrue) dirs dirs.expand(N_samples, -1) # 预测颜色和密度 rgb, sigma model(pts, dirs) # 计算透明度 delta t_vals[...,1:] - t_vals[...,:-1] delta torch.cat([delta, torch.tensor([1e10]).expand(delta[...,:1].shape)], dim-1) alpha 1 - torch.exp(-sigma.squeeze() * delta) # 累积透明度 weights alpha * torch.cumprod(1.-alpha 1e-10, dim-1) # 计算最终颜色 rgb_map torch.sum(weights[...,None] * rgb, dim-2) return rgb_map def train(model, optimizer, images, poses, hwf, N_samples64, N_rand1024): 训练循环 H, W, focal hwf for epoch in range(100000): # 随机选择一条射线 img_i np.random.randint(0, len(images)) pose poses[img_i] # 随机选择像素 i np.random.randint(0, H, sizeN_rand) j np.random.randint(0, W, sizeN_rand) # 生成射线 rays_o, rays_d get_rays(H, W, focal, pose, i, j) # 渲染 rgb_pred render_rays(model, rays_o, rays_d, near2., far6., N_samplesN_samples) # 计算损失 target images[img_i][i, j] loss F.mse_loss(rgb_pred, target) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step()在实际训练中我们还需要实现多层级采样策略即先用粗采样估计场景结构再在重要区域进行精细采样。这可以显著提高渲染质量。4. 实战技巧与性能优化训练NeRF模型时有几个关键技巧可以显著提高训练效率和最终效果GPU内存优化NeRF的体渲染过程需要处理大量射线和采样点容易导致GPU内存不足。可以通过以下方式优化分块处理将射线分成小块进行处理chunk_size 1024 * 16 # 根据GPU内存调整 for i in range(0, rays_o.shape[0], chunk_size): rgb_chunk render_rays(model, rays_o[i:ichunk_size], rays_d[i:ichunk_size], near, far, N_samples)混合精度训练使用PyTorch的自动混合精度scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): rgb_pred render_rays(model, rays_o, rays_d, near, far, N_samples) loss F.mse_loss(rgb_pred, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()训练加速技巧技巧实现方式效果提升学习率预热前1000步线性增加学习率训练更稳定权重衰减Adam优化器设置weight_decay1e-5防止过拟合学习率调度每10k步学习率减半更好收敛早停策略验证集PSNR不再提升时停止节省时间常见问题解决方案模型收敛慢检查位置编码是否正确实现确保相机参数处理正确尝试增加网络容量渲染结果模糊增加采样点数量检查位置编码的频带数量确保视角方向输入正确训练不稳定降低学习率添加梯度裁剪使用学习率预热5. 结果可视化与模型评估训练完成后我们需要评估模型的质量并可视化结果。NeRF论文中常用的评估指标包括PSNR、SSIM和LPIPS。渲染新视角def render_path(model, poses, hwf, save_path): 渲染相机路径上的所有图像 H, W, focal hwf images [] for pose in poses: rays_o, rays_d get_rays(H, W, focal, pose) rgb render_rays(model, rays_o, rays_d, near2., far6., N_samples128) images.append((rgb.cpu().numpy() * 255).astype(np.uint8)) # 保存为视频 imageio.mimsave(save_path, images, fps30)评估指标计算def compute_metrics(gt_images, pred_images): 计算PSNR和SSIM psnrs [] ssims [] for gt, pred in zip(gt_images, pred_images): # 计算PSNR mse np.mean((gt - pred) ** 2) psnr -10. * np.log(mse) / np.log(10) psnrs.append(psnr) # 计算SSIM ssim structural_similarity(gt, pred, multichannelTrue) ssims.append(ssim) return np.mean(psnrs), np.mean(ssims)在乐高数据集上一个训练良好的NeRF模型应该能达到以下指标指标粗网络精细网络PSNR28.5 dB31.2 dBSSIM0.920.94训练时间8小时15小时6. 进阶技巧与扩展应用掌握了基础NeRF实现后你可以尝试以下进阶技巧来提升模型性能动态场景处理通过添加时间维度NeRF可以用于动态场景建模。这需要修改网络结构以接受时间输入class DynamicNeRF(nn.Module): def __init__(self): super().__init__() self.time_encoder PositionalEncoding(L4) # 其余部分与标准NeRF类似大规模场景优化对于大场景可以使用以下策略空间哈希编码加速训练场景分块处理渐进式训练策略实时渲染技术NeRF的传统渲染速度较慢但可以通过以下方法加速网络蒸馏为轻量级模型预计算辐射场数据使用专用推理引擎在实际项目中我发现最耗时的部分往往是数据准备和参数调试。特别是相机参数的准确性对最终结果影响极大建议在开始训练前仔细验证数据质量。另一个实用技巧是在训练初期使用低分辨率图像待模型初步收敛后再切换到高分辨率这可以显著节省训练时间。