This repository is a PyTorchDistributedDataParallel (DDP) re-implementation of the CVPR 2022 paper CrossPoint.
-
First, the re-implementation aims to accelerate the process of training and inference by the PyTorch DDP mechanism since the original implementation by the author is for single-GPU learning and the procedure is much slower, especially when contrastive pre-training on ShapeNetRender.
-
Second, the re-implementation adds necessary comments, removes unused packages, optimizes the code style, re-organizes README, etc.
- Ubuntu 18.04
- Python 3.7.13
- PyTorch 1.11.0
- CUDA 10.2
- torchvision 0.12.0
- wandb 0.12.11
- pueue & pueued 2.0.4
conda create -n crosspoint python=3.7.13
codna activate crosspoint
pip install torch==1.11.0+cu102 torchvision==0.12.0+cu102 --extra-index-url https://download.pytorch.org/whl/cu102
pip install -r requirements.txt
pueue
is a shell command management software, we use it for scheduling the model training & inference tasks, please refer to the official page for installation and basic usage. We recommend this tool because you can run the experiments at scale with its help thus save your time.
We track the model training and fine-tuning with W&B tools. The official W&B tools may be slow and unstable since they are on remote servers, we install the local version by running the following command.
sudo groupadd docker # create `docker` group on Ubuntu
sudo usermod -aG docker <username> # add <username> to `docker` group to use docker, replace <username> with yours
docker run --rm -d -v wandb:/vol -p 28282:8080 --name wandb-local wandb/local:0.9.41
If you do not have Docker installed on your computer before, referring to the official document to finish Docker installation on Ubuntu.
CrossPoint pretrained models with DGCNN feature extractor are available here.
- Datasets are available here. Run the command below to download all the datasets (ShapeNetRender, ModelNet40, ScanObjectNN, ShapeNetPart) to reproduce the results.
cd data source download_data.sh
- The directories of the downloaded datasets should be organized as follows.
The
|- CrossPoint-DDP |---- data |-------- ShapeNetRendering |-------- modelnet40_ply_hdf5_2048 |-------- ScanObjectNN |-------- shapenet_part_seg_hdf5_data
data
directory is at the same level withmodels
,scripts
, etc.
./scripts/pt-cls.sh
./scripts/pt-partseg.sh
Run eval_ssl.ipynb
notebook to perform linear SVM object classification in both ModelNet40 and ScanObjectNN datasets.
Refer scripts/fsl_script.sh
to perform few-shot object classification.
./scripts/ft-mn.sh # finetune on ModelNet40
./scripts/ft-so.sh # finetune on ScanObjectNN
./scripts/ft-partseg.sh
The re-implementation is partially inspired by the following projects, thanks to their hard work.