Skip to content

MCG-NJU/SparseOcc

Repository files navigation

SparseOcc

This is the official PyTorch implementation for our paper:

SparseOcc: Fully Sparse 3D Occupancy Prediction
🏫 Presented by Nanjing University and Shanghai AI Lab
📧 Primary contact: Haisong Liu (afterthat97@gmail.com)
🏆 CVPR 2024 Autonomous Driving Challenge - Occupancy and Flow
📖 第三方中文解读: 自动驾驶之心AIming。谢谢你们!

Highlights

New model🥇: SparseOcc initially reconstructs a sparse 3D representation from visual inputs and subsequently predicts semantic/instance occupancy from the 3D sparse representation by sparse queries.

New evaluation metric📈: We design a thoughtful ray-based evaluation metric, namely RayIoU, to solve the inconsistency penalty along depths raised in traditional voxel-level mIoU criteria.

News

Model Zoo

Setting Pretrain Training Cost RayIoU RayPQ FPS Weights
r50_nuimg_704x256_8f nuImg 1d4h, ~12GB Memory 35.0 - 17.3 gdrive
r50_nuimg_704x256_8f_pano nuImg 1d4h, ~12GB Memory 34.5 14.0 17.3 gdrive
  • FPS is measured with Intel(R) Xeon(R) Platinum 8369B CPU and NVIDIA A100-SXM4-80GB GPU (PyTorch fp32 backend, including data loading).
  • We will release more settings in the future.

Environment

The requirements are the same as those of SparseBEV.

Install PyTorch 2.0 + CUDA 11.8:

conda create -n sparseocc python=3.8
conda activate sparseocc
conda install pytorch==2.0.0 torchvision==0.15.0 pytorch-cuda=11.8 -c pytorch -c nvidia

or PyTorch 1.10.2 + CUDA 10.2 for older GPUs:

conda create -n sparseocc python=3.8
conda activate sparseocc
conda install pytorch==1.10.2 torchvision==0.11.3 cudatoolkit=10.2 -c pytorch

Install other dependencies:

pip install openmim
mim install mmcv-full==1.6.0
mim install mmdet==2.28.2
mim install mmsegmentation==0.30.0
mim install mmdet3d==1.0.0rc6
pip install setuptools==59.5.0
pip install numpy==1.23.5

Install turbojpeg and pillow-simd to speed up data loading (optional but important):

sudo apt-get update
sudo apt-get install -y libturbojpeg
pip install pyturbojpeg
pip uninstall pillow
pip install pillow-simd==9.0.0.post1

Compile CUDA extensions:

cd models/csrc
python setup.py build_ext --inplace

Prepare Dataset

The first two steps are the same as those of SparseBEV.

  1. Download nuScenes from https://www.nuscenes.org/nuscenes, put it to data/nuscenes and preprocess it with mmdetection3d.

  2. Download the generated info file from gdrive and unzip it. These *.pkl files can also be generated with our script: gen_sweep_info.py.

  3. Download Occ3D-nuScenes occupancy GT from gdrive, unzip it, and save it to data/nuscenes/occ3d.

  4. Folder structure:

data/nuscenes
├── maps
├── nuscenes_infos_test_sweep.pkl
├── nuscenes_infos_train_sweep.pkl
├── nuscenes_infos_val_sweep.pkl
├── samples
├── sweeps
├── v1.0-test
└── v1.0-trainval
└── occ3d
    ├── scene-0001
    │   ├── 0037a705a2e04559b1bba6c01beca1cf
    │   │   └── labels.npz
    │   ├── 026155aa1c554e2f87914ec9ba80acae
    │   │   └── labels.npz
    ...
  1. (Optional) Generate the panoptic occupancy ground truth with gen_instance_info.py. The panoptic version of Occ3D will be saved to data/nuscenes/occ3d_panoptic.

Training

Train SparseOcc with 8 GPUs:

torchrun --nproc_per_node 8 train.py --config configs/sparseocc_r50_nuimg_704x256_8f.py

Train SparseOcc with 4 GPUs (i.e the last four GPUs):

export CUDA_VISIBLE_DEVICES=4,5,6,7
torchrun --nproc_per_node 4 train.py --config configs/sparseocc_r50_nuimg_704x256_8f.py

The batch size for each GPU will be scaled automatically. So there is no need to modify the batch_size in config files.

Evaluation

Single-GPU evaluation:

export CUDA_VISIBLE_DEVICES=0
python val.py --config configs/sparseocc_r50_nuimg_704x256_8f.py --weights checkpoints/sparseocc_r50_nuimg_704x256_8f.pth

Multi-GPU evaluation:

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
torchrun --nproc_per_node 8 val.py --config configs/sparseocc_r50_nuimg_704x256_8f.py --weights checkpoints/sparseocc_r50_nuimg_704x256_8f.pth

Standalone Evaluation

If you want to evaluate your own model using RayIoU, please follow the steps below:

  1. Save the predictions (shape=[200x200x16], dtype=np.uint8) with the compressed npz format. For example:
save_path = os.path.join(save_dir, sample_token + '.npz')
np.savez_compressed(save_path, pred=sem_pred)
  1. The filename for each sample is sample_token.npz, for example:
prediction/your_model
├── 000681a060c04755a1537cf83b53ba57.npz
├── 000868a72138448191b4092f75ed7776.npz
├── 0017c2623c914571a1ff2a37f034ffd7.npz
├── ...
  1. Run ray_metrics.py to evaluate on the RayIoU:
python ray_metrics.py --pred-dir prediction/your_model

Timing

FPS is measured with a single GPU:

export CUDA_VISIBLE_DEVICES=0
python timing.py --config configs/sparseocc_r50_nuimg_704x256_8f.py --weights checkpoints/sparseocc_r50_nuimg_704x256_8f.pth

Acknowledgements

Many thanks to these excellent open-source projects: