Mamba3D

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

• [2024/8] After optimizing the code and the model, Mamba3D can now achieve an overall accuracy of 92.05% on the ScanObjectNN(PB_T50_RS) dataset! We have also updated the results in the paper here.

• [2024/8] We release the pretrained weights here!

• [2024/8] We release the training and evaluation code! Pretrained weights are coming soon!

• [2024/7] Our MiniGPT-3D is also accepted by ACM MM24! We outperform existing large point-language models, using just about 1 day on 1 RTX 3090! Check it out!

• [2024/7] Ours Mamba3D is accepted by ACM MM24!

• [2024/4] We present Mamba3D, a state space model tailored for point cloud learning.

📋 TODO

• Release the training and evaluation code
• Release the pretrained weights
• Release the toy code on Colab

🎒 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.

Dataset	Pretrain	Acc	Weight
ShapeNet	Point-MAE		ckpt
ModelNet40	no	93.4	ckpt
ModelNet40	Point-MAE	94.7	ckpt
ScanObjectNN-hardest	no	91.81	ckpt
ScanObjectNN-hardest	Point-MAE	92.05	ckpt

🥧 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. Finetuning

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}
}