十亿级 3D Gaussian Splatting 的 Out-of-Core 训练：TideGS 核心技术解析

一句话总结

通过 SSD→CPU→GPU 三级存储层次 + 异步流水线 + 差分流式传输，TideGS 在单张 24GB GPU 上训练超过 10 亿个 Gaussian 原语，较标准内存方案提升 90 倍规模。

为什么需要这个？

内存墙：3DGS 的致命瓶颈

每个 3D Gaussian 原语携带的参数量很大：

参数	维度	字节数
位置 `μ`	3	12 B
协方差（四元数+缩放）	7	28 B
不透明度 `α`	1	4 B
颜色 SH 系数（3 阶）	48	192 B
合计	—	~236 B

10 亿个 Gaussian → 约 220 GB 参数存储，加上 Adam optimizer states（需要 2x 参数量）= 约 660 GB。单张 RTX 4090 的 24 GB 显存只够放 约 1100 万个 Gaussian。

稀疏性：被忽视的关键观察

传统方法把所有参数常驻 GPU 显存，但 3DGS 训练有个重要特性：

每次迭代只有当前相机视锥内可见的 Gaussian 参与计算。

对于大型城市场景，单帧可见的 Gaussian 通常只占总量的 5–15%。这意味着 85% 的参数在每次迭代中处于”睡眠”状态——为什么要让它们占据宝贵显存？

核心原理

直觉：GPU 显存变成”L3 缓存”

TideGS 的核心思路可以用计算机体系结构的缓存层次类比理解：

传统方案：
SSD → CPU RAM → GPU VRAM（所有参数常驻）
                 ^^^^^^^^ 内存墙在这里

TideGS：
SSD → CPU RAM → GPU VRAM（仅当前工作集）
↑___________↑___________↑
    后台预取         差分流式传输

GPU 显存不再是”全量参数仓库”，而是基于相机轨迹动态换入换出的工作集缓存（Working Set Cache）。

硬件层面：为什么可行？

PCIe 4.0 NVMe SSD 顺序读：约 7 GB/s
PCIe 4.0 x16 Host→Device 传输：约 32 GB/s
GPU Compute Engine 与 Copy Engine 独立：两者可真正并行执行

3DGS 的 rasterization kernel 是访存密集型，不是计算密集型。这天然存在 I/O 与计算的时间窗口可以重叠利用。

三个核心技术

Block-Virtualized Geometry：场景空间划分为规则 3D block，参数按空间局部性连续存储于 SSD，将随机 I/O 转化为顺序 I/O
Hierarchical Async Pipeline：迭代 N 的计算与迭代 N+1 所需数据的 SSD→CPU→GPU 传输完全重叠
Trajectory-Adaptive Differential Streaming：相邻迭代相机位置相近，只传输工作集的增量部分，而非每次全量传输

代码实现

Baseline：朴素的全量显存方案

import torch

class NaiveGaussianModel:
    """朴素方案：所有参数常驻 GPU 显存"""
    
    def __init__(self, n_gaussians: int):
        # 每个 Gaussian 约 236 bytes，10 亿个需要 ~220 GB → 直接 OOM
        self.positions  = torch.zeros(n_gaussians, 3,  device='cuda', requires_grad=True)
        self.quaternions= torch.zeros(n_gaussians, 4,  device='cuda', requires_grad=True)
        self.scales     = torch.zeros(n_gaussians, 3,  device='cuda', requires_grad=True)
        self.opacities  = torch.zeros(n_gaussians, 1,  device='cuda', requires_grad=True)
        self.sh_coeffs  = torch.zeros(n_gaussians, 48, device='cuda', requires_grad=True)
        
        # Adam optimizer states = 2x 参数量 → 总显存占用 3x 参数量
        self.optimizer = torch.optim.Adam([
            self.positions, self.quaternions, self.scales,
            self.opacities, self.sh_coeffs
        ], lr=1e-3)
    
    def forward(self, camera):
        # 即使大部分不可见，也要对全量 Gaussian 做视锥剔除检查 → O(N) per frame
        visible_mask = frustum_culling(self.positions, camera)
        return render_gaussians(self.positions[visible_mask], ...)

性能分析：RTX 4090（24 GB）最多支持约 11M Gaussians，超过即 OOM，无法训练更大场景。

优化 v1：块虚拟化参数管理

import numpy as np
from pathlib import Path

class BlockVirtualizedStorage:
    """
    将场景空间划分为规则 blocks，参数按空间局部性存储于 SSD
    关键：Morton code 编码使空间邻近 blocks 在文件系统中连续存放
    """
    
    def __init__(self, scene_bbox: tuple, block_size: float, ssd_path: str):
        self.block_size = block_size
        self.ssd_path = Path(ssd_path)
        bbox_min, bbox_max = scene_bbox
        self.grid_dims = np.ceil((bbox_max - bbox_min) / block_size).astype(int)
    
    def get_blocks_in_frustum(self, camera_frustum) -> list:
        """返回与视锥相交的 block ID 列表，通常只占总 blocks 的 5-15%"""
        candidate_blocks = []
        for block_id in range(int(np.prod(self.grid_dims))):
            block_bbox = self._block_id_to_bbox(block_id)
            if frustum_aabb_intersect(camera_frustum, block_bbox):
                candidate_blocks.append(block_id)
        return candidate_blocks
    
    def load_blocks(self, block_ids: list) -> dict:
        """从 SSD 批量读取，排序后接近顺序访问（比随机读快约 10x）"""
        params = {}
        for bid in sorted(block_ids):  # 排序 → 顺序读
            block_file = self.ssd_path / f"block_{bid:08d}.bin"
            params[bid] = np.fromfile(block_file, dtype=np.float32)
        return params

Morton code 编码将空间坐标交织编码为一维索引，使邻近 block 的文件偏移相近，SSD 读取速度从约 1 GB/s 提升到约 6 GB/s。

优化 v2：层次异步流水线（CUDA 多流并行）

// 核心：GPU Copy Engine 与 Compute Engine 独立，可真正并行
void training_iteration_async(
    TrainingState* state,
    const CameraInfo& current_cam,
    const CameraInfo& next_cam   // 提前预取下一帧所需数据
) {
    cudaStream_t compute_stream, transfer_stream;
    cudaStreamCreate(&compute_stream);
    cudaStreamCreate(&transfer_stream);
    
    // === 阶段1：当前帧 forward（compute_stream）===
    // 使用上一轮已预取好的数据，无需等待 I/O
    launch_gaussian_forward<<<grid, block, 0, compute_stream>>>(
        state->active_params_gpu, state->framebuffer
    );
    
    // === 阶段2：同时预取下一帧增量数据（transfer_stream，与计算并行）===
    auto next_blocks = state->storage->get_blocks_in_frustum(next_cam);
    auto new_blocks  = set_difference(next_blocks, state->active_blocks);
    
    for (int bid : new_blocks) {
        // Pinned memory → GPU DMA 直传，不占用 CPU，与 compute_stream 完全并行
        cudaMemcpyAsync(
            state->gpu_buffer + bid * BLOCK_PARAM_SIZE,
            state->cpu_pinned_buffer + bid * BLOCK_PARAM_SIZE,
            BLOCK_PARAM_SIZE * sizeof(float),
            cudaMemcpyHostToDevice, transfer_stream
        );
    }
    
    // Backward 必须等 Forward 完成（有数据依赖）
    cudaStreamSynchronize(compute_stream);
    launch_gaussian_backward<<<grid, block, 0, compute_stream>>>(
        state->active_params_gpu, state->grad_buffer
    );
    
    cudaStreamSynchronize(transfer_stream);
    state->swap_active_blocks(new_blocks, get_lru_blocks(state));
}

为什么更快：I/O 等待时间从 35% 下降到 <5%，等效训练吞吐量提升约 1.8x。

优化 v3：轨迹自适应差分流式传输

class DifferentialStreamingManager:
    """
    只传增量，不传全集
    利用相机轨迹连续性：相邻帧工作集重叠度高（通常 85-90%）
    """
    
    def __init__(self, gpu_cache_capacity_blocks: int):
        self.capacity = gpu_cache_capacity_blocks
        self.active_blocks: set = set()
        self._stats = {"total_naive": 0, "actual_transferred": 0}
    
    def compute_delta(self, next_active: set) -> tuple[set, set]:
        """差分计算：返回（需加载的新 blocks，需驱逐的旧 blocks）"""
        to_load  = next_active - self.active_blocks   # 新进入视锥
        to_evict = self.active_blocks - next_active   # 离开视锥
        
        self._stats["total_naive"]        += len(next_active)
        self._stats["actual_transferred"] += len(to_load)
        return to_load, to_evict
    
    def apply_delta(self, to_load: set, to_evict: set, block_params: dict):
        for bid in to_evict:
            self._writeback_gradients_to_cpu(bid)  # 写回梯度，不丢更新
            self.active_blocks.remove(bid)
        for bid in to_load:
            self._upload_to_gpu(block_params[bid])
            self.active_blocks.add(bid)
    
    @property
    def transfer_ratio(self) -> float:
        """实际传输量 / 朴素全量传输量，越低越好"""
        if self._stats["total_naive"] == 0:
            return 1.0
        return self._stats["actual_transferred"] / self._stats["total_naive"]

实测：在连续视频序列上，差分传输率约为全量传输的 12–18%，等效带宽利用率提升 5–8x。

常见错误

错误1：用普通内存代替 Pinned Memory

# 错误：pageable memory，每次 DMA 前 CPU 需要额外拷贝一次
staging = np.zeros((N, PARAM_DIM), dtype=np.float32)

# 正确：pinned memory，允许 DMA 直传，传输速度提升约 2x
staging = torch.zeros(N, PARAM_DIM, dtype=torch.float32).pin_memory()

错误2：隐式同步打断流水线

# 错误：.item() / .numpy() 会隐式调用 cudaDeviceSynchronize()，流水线断流
loss_val = loss.item()   # GPU 停下来等 CPU 读数据

# 正确：保持在 tensor 上操作，延迟到真正需要时再同步
loss_val = loss.detach()  # 不触发同步

性能实测

测试环境：RTX 4090 (24 GB VRAM), PCIe 4.0 NVMe SSD, CUDA 12.1

实现版本	最大 Gaussian 数量	训练速度 (iter/s)	大场景 PSNR
标准 3DGS	~11M	2.1	24.3 dB
朴素 Out-of-Core（同步）	~100M	0.3	25.8 dB
+ 块虚拟化	~300M	0.8	26.5 dB
+ 异步流水线	~700M	1.4	27.1 dB
TideGS 完整版	>1B	1.7	27.6 dB

什么时候用 / 不用？

适用场景	不适用场景
城市级大场景重建（室外、无界场景）	小场景（房间、单物体），11M 已足够
有高质量 PCIe 4.0+ NVMe SSD	机械硬盘（HDD）环境，顺序读仅 200 MB/s
需要超高精度渲染（飞行漫游、影视级）	实时推理应用（与训练阶段无关）
相机轨迹连续（差分流式有效）	随机跳跃式相机采样（差分优化失效）

调试技巧

1. 验证流水线重叠是否生效

用 Nsight Systems 查看 CUDA 流时序图。理想状态应看到 cudaMemcpyAsync 和 kernel 执行的时间条完全重叠。如果是串行，检查是否意外插入了 cudaDeviceSynchronize() 或 .item() 调用。

2. Block 大小调优

Block 太小 → 元数据开销大，I/O 请求数过多；Block 太大 → 每次载入冗余数据，显存浪费。对于室外大场景，8m×8m×8m 通常是合理起点，可根据场景密度调整。

3. LRU 驱逐策略的抖动陷阱

相机来回移动时，简单 LRU 会产生缓存抖动（cache thrashing）：同一 block 被反复驱逐和加载。改进方案：加入基于轨迹预测的保留策略，若未来若干帧内相机将返回该区域，保留对应 blocks 不驱逐。

4. 梯度写回的一致性

驱逐 block 时必须将 GPU 上已积累的梯度写回 CPU，否则那部分参数在下次载入时的 optimizer state 会丢失，导致训练不收敛。这是实现中最容易忘记的细节。