CauchyFood / Point-MAE

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool


Masked Autoencoders for Point Cloud Self-supervised Learning, arxiv

In this work, we present a novel scheme of masked autoencoders for point cloud self-supervised learning, termed as Point-MAE. Our Point-MAE is neat and efficient, with minimal modifications based on the properties of the point cloud. In classification tasks, Point-MAE outperforms all the other self-supervised learning methods on ScanObjectNN and ModelNet40. Point-MAE also advances state-of-the-art accuracies by 1.5%-2.3% in the few-shot learning on ModelNet40.

1. Requirements

PyTorch >= 1.7.0; python >= 3.7; CUDA >= 9.0; GCC >= 4.9; torchvision;

pip install -r requirements.txt
# Chamfer Distance & emd
cd ./extensions/chamfer_dist
python install --user
cd ./extensions/emd
python install --user
# PointNet++
pip install "git+git://"
pip install --upgrade

2. Datasets

We use ShapeNet, ScanObjectNN, ModelNet40 and ShapeNetPart in this work. See for details.

3. Point-MAE Models

Task Dataset Config Acc. Download
Pre-training ShapeNet pretrain.yaml N.A. here
Classification ScanObjectNN finetune_scan_hardest.yaml 84.52% here
Classification ScanObjectNN finetune_scan_objbg.yaml 88.29% here
Classification ScanObjectNN finetune_scan_objonly.yaml 90.01% here
Classification ModelNet40(1k) finetune_modelnet.yaml 93.80% here
Classification ModelNet40(8k) finetune_modelnet_8k.yaml 94.04% here
Part segmentation ShapeNetPart segmentation 86.1% mIoU here
Task Dataset Config 5w10s Acc. (%) 5w20s Acc. (%) 10w10s Acc. (%) 10w20s Acc. (%)
Few-shot learning ModelNet40 fewshot.yaml 96.3 ± 2.5 97.8 ± 1.8 92.6 ± 4.1 95.0 ± 3.0

4. Point-MAE Pre-training

To pretrain Point-MAE on ShapeNet training set, run the following command. If you want to try different models or masking ratios etc., first create a new config file, and pass its path to --config.

CUDA_VISIBLE_DEVICES=<GPU> python --config cfgs/pretrain.yaml --exp_name <output_file_name>

5. Point-MAE Fine-tuning

Fine-tuning on ScanObjectNN, run:

CUDA_VISIBLE_DEVICES=<GPUs> python --config cfgs/finetune_scan_hardest.yaml \
--finetune_model --exp_name <output_file_name> --ckpts <path/to/pre-trained/model>

Fine-tuning on ModelNet40, run:

CUDA_VISIBLE_DEVICES=<GPUs> python --config cfgs/finetune_modelnet.yaml \
--finetune_model --exp_name <output_file_name> --ckpts <path/to/pre-trained/model>

Voting on ModelNet40, run:

CUDA_VISIBLE_DEVICES=<GPUs> python --test --config cfgs/finetune_modelnet.yaml \
--exp_name <output_file_name> --ckpts <path/to/best/fine-tuned/model>

Few-shot learning, run:

CUDA_VISIBLE_DEVICES=<GPUs> python --config cfgs/fewshot.yaml --finetune_model \
--ckpts <path/to/pre-trained/model> --exp_name <output_file_name> --way <5 or 10> --shot <10 or 20> --fold <0-9>

Part segmentation on ShapeNetPart, run:

cd segmentation
python --ckpts <path/to/pre-trained/model> --root path/to/data --learning_rate 0.0002 --epoch 300

6. Visualization

Visulization of pre-trained model on ShapeNet validation set, run:

python --test --ckpts <path/to/pre-trained/model> --config cfgs/pretrain.yaml --exp_name <name>


Our codes are built upon Point-BERT, Pointnet2_PyTorch and Pointnet_Pointnet2_pytorch


    title={Masked Autoencoders for Point Cloud Self-supervised Learning},
    author={Yatian Pang and Wenxiao Wang and Francis E. H. Tay and Wei Liu and Yonghong Tian and Li Yuan},



Language:Python 90.2%Language:Cuda 8.8%Language:C++ 1.0%