Official Pytorch implementation of the method NCLR. More details can be found in the paper:
Self-supervised Learning of LiDAR 3D PointClouds via 2D-3D Neural Calibration, Arxiv 2024 [arXiv] by Yifan Zhang, Siyu Ren, Junhui Hou, Jinjian Wu, Yixuan Yuan, Guangming Shi
If you use NCLR in your research, please cite:
@article{zhang2024nclr,
title={Self-supervised Learning of LiDAR 3D Point Clouds via 2D-3D Neural Calibration},
author={Zhang, Yifan and Ren, Siyu and Hou, Junhui and Wu, Jinjian and Yuan, Yixuan and Shi, Guangming},
journal={arXiv preprint arXiv:2401.12452},
year={2024}
}
To install the various dependencies, you can run pip install -r requirements.txt
.
The code provided can be used with nuScenes, SemanticKITTI, and SemanticPOSS. Put the datasets you intend to use in the datasets folder (a symbolic link is accepted).
SR-UNet pre-trained on nuScenes
SR-UNet pre-trained on SemanticKITTI
VoxelNet pre-trained on nuScenes
When using MinkowskiEngine (on SemanticKITTI), please set the OMP_NUM_THREADS environment variable to your number of CPU cores
Config file for SemanticKITTI is included for MinkowskiEngine by default to keep retro-compatibility with previous work, while for nuScenes it uses Torchsparse which is generally faster. Switching between libraries in the config files is easy. While architectures are similar, weights from one library cannot easily be transferred to the other.
- On nuScenes:
python train.py --config_file cfgs/pretrain_ns_minkunet.yaml --name minkunet_nclr_ns
- On SemanticKITTI:
python train.py --config_file cfgs/pretrain_sk_minkunet.yaml --name minkunet_nclr_sk
The specific code for downstream semantic segmentation has been adapted from ALSO.
Method | 0.1% | 1% | 10% | 50% | 100% |
---|---|---|---|---|---|
Random init. | 21.6 | 35.0 | 57.3 | 69.0 | 71.2 |
PointContrast | 27.1 | 37.0 | 58.9 | 69.4 | 71.1 |
DepthContrast | 21.7 | 34.6 | 57.4 | 69.2 | 71.2 |
ALSO | 26.2 | 37.4 | 59.0 | 69.8 | 71.8 |
NCLR | 26.6 | 37.8 | 59.5 | 71.2 | 72.7 |
To launch a downstream experiment, with a Torchsparse SR-Unet, you can use these commands in addition with cfg.downstream.checkpoint_dir=[checkpoint directory] cfg.downstream.checkpoint_name=[checkpoint name]
cd downstream
# 100%
python train_downstream_semseg.py cfg=nuscenes_torchsparse cfg.downstream.max_epochs=30 cfg.downstream.val_interval=5 cfg.downstream.skip_ratio=1
# 50%
python train_downstream_semseg.py cfg=nuscenes_torchsparse cfg.downstream.max_epochs=50 cfg.downstream.val_interval=5 cfg.downstream.skip_ratio=2
# 10%
python train_downstream_semseg.py cfg=nuscenes_torchsparse cfg.downstream.max_epochs=100 cfg.downstream.val_interval=10 cfg.downstream.skip_ratio=10
# 1%
python train_downstream_semseg.py cfg=nuscenes_torchsparse cfg.downstream.max_epochs=500 cfg.downstream.val_interval=50 cfg.downstream.skip_ratio=100
# 0.1%
python train_downstream_semseg.py cfg=nuscenes_torchsparse cfg.downstream.max_epochs=1000 cfg.downstream.val_interval=100 cfg.downstream.skip_ratio=1000
To launch a downstream experiment, with a Minkowski SR-Unet, you can use these commands in addition with cfg.downstream.checkpoint_dir=[checkpoint directory] cfg.downstream.checkpoint_name=[checkpoint name]
cd downstream
# 100%
python train_downstream_semseg.py cfg=nuscenes_minkowski cfg.downstream.max_epochs=30 cfg.downstream.val_interval=5 cfg.downstream.skip_ratio=1
# 50%
python train_downstream_semseg.py cfg=nuscenes_minkowski cfg.downstream.max_epochs=50 cfg.downstream.val_interval=5 cfg.downstream.skip_ratio=2
# 10%
python train_downstream_semseg.py cfg=nuscenes_minkowski cfg.downstream.max_epochs=100 cfg.downstream.val_interval=10 cfg.downstream.skip_ratio=10
# 1%
python train_downstream_semseg.py cfg=nuscenes_minkowski cfg.downstream.max_epochs=500 cfg.downstream.val_interval=50 cfg.downstream.skip_ratio=100
# 0.1%
python train_downstream_semseg.py cfg=nuscenes_minkowski cfg.downstream.max_epochs=1000 cfg.downstream.val_interval=100 cfg.downstream.skip_ratio=1000
python train.py --config_file cfgs/pretrain_ns_spconv.yaml --name voxelnet_nclr_ns
Please use the code of OpenPCDet with default parameters for SECOND or PVRCNN and with no multiprocessing to retain compatibility with previous work and this one.
Panoptic segmentation baseline MinkowskiPanoptic implemented based on the MinkowskiEngine library
Part of the codebase has been adapted from OpenPCDet, ALSO, valeoai, and SLidR.
- Initial release.
- Add license. See here for more details.
- Add installation details.
- Add data preparation details.
- Add evaluation details.
- Add training details.
- LTA
- Add pre-trained weights.
Here are some of the methods I recommend for 3D representation learning:
NCLR is released under the Apache 2.0 license.