GaussianSTORM

Training is performed using main_storm.py. The following example demonstrates a multi-GPU setup to reproduce the STORM-B/8 model.

Key Configuration Details:

• batch_size is per-GPU. The global batch size is calculated as batch_size × #GPUs × #nodes.
• Checkpoints and logs are saved to work_dirs/<project>/<exp_name>.
• For a full list of arguments, refer to main_storm.py.

# Multi-GPU training example
torchrun --nproc_per_node=8 main_storm.py \
    --project 0504_storm \
    --exp_name 0504_pixel_storm \
    --data_root ../storm2.3/data/STORM2 \ # replace this with your data root.
    --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

After downloading the raw .tfrecord files, use preprocess.py to extract and organize components like images, LiDAR, and calibration data into a processed directory structure.

Arguments:

• --data_root: Path to the raw Waymo data.
• --target_dir: Path where processed data should be saved.
• --dataset: Set to waymo.
• --split: Either training or validation.
• --scene_list_file: Path to the split file containing scene names.
• --scene_ids: (Optional) Specific scene IDs to process.
• --num_workers: Number of worker threads.
• --process_keys: A space-separated list of components to extract (e.g., images lidar calib pose dynamic_masks ground).
• --json_folder_to_save: Path to save the resulting annotations.

# Preprocess specific scenes
python preprocess.py \
    --data_root data/waymo/raw/ \
    --target_dir data/waymo/processed \
    --dataset waymo \
    --split training \
    --scene_list_file data/dataset_scene_list/waymo_train_list.txt \
    --scene_ids 700 754 23 \
    --num_workers 8 \
    --process_keys images lidar calib pose dynamic_masks ground \
    --json_folder_to_save data/STORM_data/annotations/waymo 

Use the preproc/waymo_download.py script to download specific scenes or entire splits from the Waymo Open Dataset. You must have a Waymo Open Dataset account and the gcloud SDK installed and authenticated.

Arguments:

• --target_dir: The directory where raw data will be stored (e.g., ./data/waymo/raw/training).
• --split_file: Path to the text file containing the scene list (e.g., data/dataset_scene_list/waymo_train_list.txt).
• --scene_ids: (Optional) A space-separated list of specific scene IDs to download. If omitted, all scenes in the --split_file are downloaded.

# Download specific scenes
python preproc/waymo_download.py \
    --target_dir ./data/waymo/raw/training \
    --split_file data/dataset_scene_list/waymo_train_list.txt \
    --scene_ids 700 754 23

# Download all scenes in the training split
python preproc/waymo_download.py \
    --target_dir ./data/waymo/raw/training \
    --split_file data/dataset_scene_list/waymo_train_list.txt

Follow these steps to set up the STORM environment. The project was tested with CUDA 12.1, PyTorch 2.3, and an NVIDIA A100. You may need to adjust CUDA/PyTorch versions to match your environment.

1. Clone the repository.
2. Create and activate a Conda environment with Python 3.10.
3. Install Python dependencies via requirements.txt.
4. Install gsplat for batch-wise rendering support. Note that gsplat installation can be machine-dependent. If the specific commit installation fails, try installing the latest version.

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

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

# install python dependencies
pip install -r requirements.txt

# install gsplat (for batch-wise rendering support)
pip install git+https://github.com/nerfstudio-project/gsplat.git@2b0de894232d21e8963179a7bbbd315f27c52c9c
# If the above fails, try:
# pip install git+https://github.com/nerfstudio-project/gsplat.git

The following structure is expected after successful preprocessing and sky mask extraction:

ProjectPath/data/
└── waymo/
    ├── raw/                 # Original .tfrecord files
    │    └── ...
    └── processed/
         └── [training|validation]/
              └── [scene_id]/  # e.g., 000, 001...
                 ├── cam_to_ego/         # {cam_id}.txt
                 ├── cam_to_world/       # {timestep:03d}_{cam_id}.txt
                 ├── depth_flows_4/      # {timestep:03d}_{cam_id}.npy
                 ├── dynamic_masks/      # {timestep:03d}_{cam_id}.png
                 ├── ego_to_world/       # {timestep:03d}.txt
                 ├── ground_label_4/     # {timestep:03d}.txt
                 ├── images/             # {timestep:03d}_{cam_id}.jpg
                 ├── images_4/           # {timestep:03d}_{cam_id}.jpg
                 ├── intrinsics/         # {cam_id}.txt
                 ├── lidar/               # {timestep:03d}.bin
                 └── sky_masks/          # {timestep:03d}_{cam_id}.png