PVT: Point-Voxel Transformer for 3D Deep Learning

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

Paper and Citation

The paper can be downloaded from arXiv.<BR/> 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

Pretrained Models

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