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
-
[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.
- Release the training and evaluation code
- Release the pretrained weights
- Release the toy code on Colab
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.
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 |
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
To fine-tune Mamba3D on ScanObjectNN/Modelnet40, run:
# Note: change config files for different dataset
bash script/run_finetune.sh
We would like to thank the authors of Mamba, Vision Mamba, and Point-MAE for their great works and repos.
If you have any questions or are looking for cooperation in related fields, please contact Xu Han via xhanxu@hust.edu.cn.
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}
}