GaussianSTORM开源项目教程

1. 项目介绍

GaussianSTORM是一个基于PyTorch的开源项目，它实现了STORM（Spatio-Temporal Reconstruction Model for Large-Scale Outdoor Scenes）模型。该模型是一种用于大规模户外场景的时空重建方法，能够从稀疏的多视角序列中快速、自监督地重建动态场景。它同时学习3D高斯分布和场景流，支持实时渲染和运动分割，并且在动态区域上比场景优化和其他通用模型具有更高的性能。

2. 项目快速启动

环境准备

• CUDA 12.1
• PyTorch 2.3
• NVIDIA A100（或其他兼容GPU）

克隆项目

git clone https://github.com/NVlabs/GaussianSTORM.git
cd GaussianSTORM

创建conda环境

conda create -n storm python=3.10 -y
conda activate storm

安装Python依赖

pip install -r requirements.txt

安装gsplat（用于批处理渲染支持）

pip install git+https://github.com/nerfstudio-project/gsplat.git@2b0de894232d21e8963179a7bbbd315f27c52c9c

快速实验

我们提供了一个Waymo Open Dataset的子集（3个序列）用于快速实验。

下载数据子集（约600MB）

gdown 14fapsAGoMCQ5Ky82cg2X6bk-mLQ7fdCF
tar -xf STORM_subset.tar.gz

运行单GPU推理演示

python inference.py \
--project storm_playground --exp_name visualization \
--data_root data/STORM_subset \
--model STORM-B/8 --num_motion_tokens 16 \
--use_sky_token --use_affine_token \
--load_depth --load_flow --load_ground \
--load_from $CKPT_PTH

请注意，CKPT_PTH指的是检查点路径。由于官方检查点目前无法提供，请参考问题页面获取非官方检查点。

3. 应用案例和最佳实践

数据准备

为了准备Waymo Open Dataset，请参考Waymo Data的官方指南。对于NuScenes和Argoverse2数据集，目前尚未包含准备指南，但可能会根据需求添加。

训练模型

以下是一个多GPU训练的示例，用于重现论文中的STORM-B/8模型：

torchrun --nproc_per_node=8 main_storm.py \
--project 0504_storm \
--exp_name 0504_pixel_storm \
--data_root ../storm2.3/data/STORM2
...
# 请替换为您的数据根目录
--batch_size 4 --num_iterations 100000 --lr_sched constant \
--model STORM-B/8 --num_motion_tokens 16 \
--use_sky_token --use_affine_token \
--load_depth --load_flow --load_ground \
--enable_depth_loss --enable_flow_reg_loss --flow_reg_coeff 0.005 --enable_sky_opacity_loss \
--enable_perceptual_loss --perceptual_loss_start_iter 5000 \
--enable_wandb \
--auto_resume

评估模型

torchrun --nproc_per_node=8 main_storm.py \
--project 0504_storm \
--exp_name 0504_pixel_storm \
--data_root ../storm2.3/data/STORM2
...
# 请替换为您的数据根目录
--batch_size 4 --num_iterations 100000 --lr_sched constant \
--model STORM-B/8 --num_motion_tokens 16 \
--use_sky_token --use_affine_token \
--load_depth --load_flow --load_ground \
--enable_depth_loss --enable_flow_reg_loss --flow_reg_coeff 0.005 --enable_sky_opacity_loss \
--enable_perceptual_loss --perceptual_loss_start_iter 5000 \
--auto_resume \
--evaluate

4. 典型生态项目

目前，GaussianSTORM项目主要围绕大规模户外场景的时空重建。在开源社区中，类似的项目还包括用于3D重建、场景理解和其他计算机视觉任务的项目。例如，NERF（Neural Radiance Fields）和PointNet等，都是与GaussianSTORM互补的开源项目，共同推动了计算机视觉领域的发展。