PVT: Point-Voxel Transformer for 3D Deep Learning

May 26, 2022 · View on GitHub

We will release the latest code of our PVTv4 at arXiv upon acceptance!

Paper and Citation

The paper can be downloaded from arXiv.
If you like our work and think it helpful to your project, please cite it as follows.

@article{zhang2021point,
  title={PVT: Point-Voxel Transformer for 3D Deep Learning},
  author={Zhang, Cheng and Wan, Haocheng and Liu, Shengqiang and Shen, Xinyi and Wu, Zizhao},
  journal={arXiv preprint arXiv:2108.06076},
  year={2021}
}

Prerequisites

The code is built with following libraries (see requirements.txt):

Python >= 3.7
PyTorch >= 1.3
numba
numpy
scipy
six
tensorboardX >= 1.2
plyfile
h5py
timm

Data Preparation

ModelNet40

Download alignment ModelNet40 here and save in data/modelnet40_normal_resampled/.

S3DIS

We follow the data pre-processing in PointCNN. The code for preprocessing the S3DIS dataset is located in data/s3dis/. One should first download the dataset from here, then run

python data/s3dis/prepare_data.py

ShapeNet

We follow the data pre-processing in PointNet2. Please download the dataset from hereand save in data/shapenetcore_partanno_segmentation_benchmark_v0_normal/.

KITTI

For Frustum-PointNet backbone, we follow the data pre-processing in Frustum-Pointnets. One should first download the ground truth labels from here, then run

unzip data_object_label_2.zip
mv training/label_2 data/kitti/ground_truth
./data/kitti/frustum/download.sh

Here we provide a pretrained model on ModelNet40. The accuracy might vary a little bit compared to the paper, since we re-train some of the models for reproducibility. The path of the model is in ./checkpoints/cls/model.t7

Example training and testing

#train
python main_cls.py --exp_name=cls --num_points=1024 --use_sgd=True --batch_size 32 --epochs 200 --lr 0.001

#test
python main_cls.py --exp_name=cls --num_points=1024 --use_sgd=True --eval=True --model_path=checkpoints/cls/model.t7 --test_batch_size 32