SeFlow: A Self-Supervised Scene Flow Method in Autonomous Driving

[poster comming soon] [video coming soon]

2024/07/16 17:18: Most of codes already uploaded and tested. You can to try training directly by downloading demo data or pretrained weight for evaluation. The process script will be public when the paper published.

Pre-trained weights for models are available in Zenodo link. Check usage in 2. Evaluation or 3. Visualization.

Task: Self-Supervised Scene Flow Estimation in Autonomous Driving. No human-label needed. Real-time inference (15-20Hz in RTX3090).

We directly follow our previous work code structure, so you may want to start from the easier one with supervised learning first: Try DeFlow. Then you will find this is simple to you (things about how to train under self-supervised). Here are Scripts quick view in this repo:

• dataprocess/extract_*.py : pre-process data before training to speed up the whole training time. [Dataset we included now: Argoverse 2 and Waymo. more on the way: Nuscenes, custom data.]

• 0_process.py: process data with save dufomap, cluster labels inside file. Only needed for training.

• 1_train.py: Train the model and get model checkpoints. Pls remember to check the config.

• 2_eval.py : Evaluate the model on the validation/test set. And also upload to online leaderboard.

• 3_vis.py : For visualization of the results with a video.

🎁 One repository, All methods!

You can try following methods in our code without any effort to make your own benchmark.

• SeFlow (Ours 🚀): ECCV 2024
• DeFlow (Ours 🚀): ICRA 2024
• FastFlow3d: RA-L 2021
• ZeroFlow: ICLR 2024, their pre-trained weight can covert into our format easily through the script.
• NSFP: NeurIPS 2021, faster 3x than original version because of our CUDA speed up, same (slightly better) performance. Done coding, public after review.
• FastNSF: ICCV 2023. Done coding, public after review.

• ... more on the way

💡: Want to learn how to add your own network in this structure? Check Contribute section and know more about the code. Fee free to pull request!

0. Setup

Environment: Same to DeFlow. And even lighter here with extracting mmcv module we needed into cuda assets.

git clone --recursive https://github.com/KTH-RPL/SeFlow.git
cd SeFlow && mamba env create -f environment.yaml

CUDA package (need install nvcc compiler), the compile time is around 1-5 minutes:

mamba activate seflow
# change it if you use different cuda version (I tested 11.3, 11.4, 11.7, 11.8 all works)
export PATH=/usr/local/cuda-11.7/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64:$LD_LIBRARY_PATH

cd assets/cuda/mmcv && python ./setup.py install && cd ../../..
cd assets/cuda/chamfer3D && python ./setup.py install && cd ../../..

Or you always can choose Docker which isolated environment and free yourself from installation, you can pull it by. If you have different arch, please build it by yourself cd SeFlow && docker build -t zhangkin/seflow by going through build-docker-image section.

# option 1: pull from docker hub
docker pull zhangkin/seflow

# run container
docker run -it --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name seflow zhangkin/seflow /bin/zsh

1. Run & Train

Note: Prepare raw data and process train data only needed run once for the task. No need repeat the data process steps till you delete all data. We use wandb to log the training process, and you may want to change all entity="kth-rpl" to your own entity.

Data Preparation

Check dataprocess/README.md for downloading tips for the raw Argoverse 2 dataset. Or maybe you want to have the mini processed dataset to try the code quickly, We directly provide one scene inside train and val. It already converted to .h5 format and processed with the label data. You can download it from Zenodo and extract it to the data folder. And then you can skip following steps and directly run the training script.

wget https://zenodo.org/record/12751363/files/demo_data.zip
unzip demo_data.zip -p /home/kin/data/av2

Prepare raw data

Checking more information (step for downloading raw data, storage size, #frame etc) in dataprocess/README.md. Extract all data to unified h5 format. [Runtime: Normally need 10 mins finished run following commands totally in my desktop, 45 mins for the cluster I used]

python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir /home/kin/data/av2 --output_dir /home/kin/data/av2/preprocess_v2
python dataprocess/extract_av2.py --av2_type sensor --data_mode val --mask_dir /home/kin/data/av2/3d_scene_flow
python dataprocess/extract_av2.py --av2_type sensor --data_mode test --mask_dir /home/kin/data/av2/3d_scene_flow

Process train data

Process train data for self-supervised learning. Only training data needs this step. [Runtime: Normally need 15 hours for my desktop, 3 hours for the cluster with five available nodes parallel running.]

python 0_process.py --data_dir /home/kin/data/av2/preprocess_v2/sensor/train --scene_range 0,701

Train the model

Train SeFlow needed to specify the loss function, we set the config of our best model in the leaderboard. [Runtime: Around 11 hours in 4x A100 GPUs.]

python 1_train.py model=deflow lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" "model.target.num_iters=2" "model.val_monitor=val/Dynamic/Mean"

Or you can directly download the pre-trained weight from Zenodo and skip the training step.

Other Benchmark Models

You can also train the supervised baseline model in our paper with the following command. [Runtime: Around 10 hours in 4x A100 GPUs.]

python 1_train.py model=fastflow3d lr=2e-4 epochs=20 batch_size=16 loss_fn=ff3dLoss
python 1_train.py model=deflow lr=2e-4 epochs=20 batch_size=16 loss_fn=deflowLoss

Note

You may found the different settings in the paper that is all methods are enlarge learning rate to 2e-4 and decrease the epochs to 20 for faster converge and better performance. However, we kept the setting on lr=2e-6 and 50 epochs in (SeFlow & DeFlow) paper experiments for the fair comparison with ZeroFlow where we directly use their provided weights. We suggest afterward researchers or users to use the setting here (larger lr and smaller epoch) for faster converge and better performance.

2. Evaluation

You can view Wandb dashboard for the training and evaluation results or upload result to online leaderboard.

Since in training, we save all hyper-parameters and model checkpoints, the only thing you need to do is to specify the checkpoint path. Remember to set the data path correctly also.

# downloaded pre-trained weight, or train by yourself
wget https://zenodo.org/records/12751363/files/seflow_best.ckpt

# it will directly prints all metric
python 2_eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=val

# it will output the av2_submit.zip or av2_submit_v2.zip for you to submit to leaderboard
python 2_eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=1
python 2_eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=2

And the terminal will output the command for you to submit the result to the online leaderboard. You can follow this section for evalai.

Check all detailed result files (presented in our paper Table 1) in this discussion.

3. Visualization

We provide a script to visualize the results of the model also. You can specify the checkpoint path and the data path to visualize the results. The step is quickly similar to evaluation.

python 3_vis.py checkpoint=/home/kin/seflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess_v2/sensor/vis

# The output of above command will be like:
Model: DeFlow, Checkpoint from: /home/kin/model_zoo/v2/seflow_best.ckpt
We already write the flow_est into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:
python tools/scene_flow.py --flow_mode 'seflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
Enjoy! ^v^ ------ 

# Then run the command in the terminal:
python tools/scene_flow.py --flow_mode 'seflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis

seflow.mp4

Cite & Acknowledgements

@article{zhang2024seflow,
  author={Zhang, Qingwen and Yang, Yi and Li, Peizheng and Andersson, Olov and Jensfelt, Patric},
  title={SeFlow: A Self-Supervised Scene Flow Method in Autonomous Driving},
  journal={arXiv preprint arXiv:2407.01702},
  year={2024}
}
@article{zhang2024deflow,
  author={Zhang, Qingwen and Yang, Yi and Fang, Heng and Geng, Ruoyu and Jensfelt, Patric},
  title={DeFlow: Decoder of Scene Flow Network in Autonomous Driving},
  journal={arXiv preprint arXiv:2401.16122},
  year={2024}
}

💞 Thanks to RPL member: Li Ling helps revise our SeFlow manuscript. Thanks to Kyle Vedder, who kindly opened his code (ZeroFlow) including pre-trained weights, and discussed their result with us which helped this work a lot.

This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation and Prosense (2020-02963) funded by Vinnova. The computations were enabled by the supercomputing resource Berzelius provided by National Supercomputer Centre at Linköping University and the Knut and Alice Wallenberg Foundation, Sweden.

❤️: DeFlow, BucketedSceneFlowEval