Home

Awesome

Mamba3D

PWC PWC PWC

This repository contains the official implementation of the paper:

[ACM MM 24] Mamba3D: Enhancing Local Features for 3D Point Cloud Analysis via State Space Model

๐Ÿ“ฐ News

<div style="text-align: center"> <img src="media/mamba3d_total_v2.png" /> </div>

๐Ÿ“‹ TODO

๐ŸŽ’ 1. Requirements

Tested on: PyTorch == 1.13.1; python == 3.8; CUDA == 11.7

pip install -r requirements.txt
# Chamfer Distance & emd
cd ./extensions/chamfer_dist
python setup.py install --user
cd ./extensions/emd
python setup.py install --user
# PointNet++
pip install "git+https://github.com/erikwijmans/Pointnet2_PyTorch.git#egg=pointnet2_ops&subdirectory=pointnet2_ops_lib"
# GPU kNN
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl

# Mamba install
pip install causal-conv1d==1.1.1
pip install mamba-ssm==1.1.1

More detailed settings can be found in mamba3d.yaml.

๐Ÿงพ 2. Datasets & Pretrained Weights

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

You can find the pre-trained weights here. Or, specifically as follows.

DatasetPretrainAccWeight
ShapeNetPoint-MAEckpt
ModelNet40no93.4ckpt
ModelNet40Point-MAE94.7ckpt
ScanObjectNN-hardestno91.81ckpt
ScanObjectNN-hardestPoint-MAE92.05ckpt

๐Ÿฅง 3. Training from scratch

To train Mamba3D on ScanObjectNN/Modelnet40 from scratch, run:

# Note: change config files for different dataset
bash script/run_scratch.sh

To vote on ScanObjectNN/Modelnet40, run:

# Note: change config files for different dataset
bash script/run_vote.sh

Few-shot learning, run:

bash script/run_fewshot.sh
<!-- ## ๐ŸŸ 4. Pretraining & Finetuning -->

๐ŸŸ 4. Finetuning

<!-- To pre-train Mamba3D on ShapeNet, run: ``` # Note: change config files for different dataset bash script/run_pretrain.sh # or CUDA_VISIBLE_DEVICES=<GPU> python main.py --config cfgs/pretrain.yaml --exp_name <output_file_name> ``` -->

To fine-tune Mamba3D on ScanObjectNN/Modelnet40, run:

# Note: change config files for different dataset
bash script/run_finetune.sh

๐Ÿ˜Š Acknowledgement

We would like to thank the authors of Mamba, Vision Mamba, and Point-MAE for their great works and repos.

๐Ÿ˜€ Contact

If you have any questions or are looking for cooperation in related fields, please contact Xu Han via xhanxu@hust.edu.cn.

๐Ÿ“š Citation

If you find our work helpful, please consider citing:

@article{han2024mamba3d,
  title={Mamba3D: Enhancing Local Features for 3D Point Cloud Analysis via State Space Model},
  author={Han, Xu and Tang, Yuan and Wang, Zhaoxuan and Li, Xianzhi},
  journal={arXiv preprint arXiv:2404.14966},
  year={2024}
}