Fork of official Mamba3D repository with updates to more recent versions of pytorch, CUDA, and some other packages

Changes

• Updated to CUDA 12.1 and corresponding conda pytorch install (following pytorch installation guidelines along the lines of conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia)
• Updated to python 3.11 (note 3.12 does not seem to work)
• Update to Pointnet: Fetch code from here https://github.com/evenrose/Pointnet2_PyTorch.git, then install (in directory) with pip install pointnet2_ops_lib/

Additional pip installs compared to main instructions below:

pip install causal-conv1d # Note no version
pip install openTSNE
pip install fvcore
pip install mmcv-full
pip install thop
pip install mamba-ssm

Below is the original README.

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] We've slightly optimized the architecture, such as some normalization layers and hyperparameters.

• [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.

• [2024/4] Mamba3D surpasses Transformer-based counterparts and concurrent works in multiple tasks, achieving multiple SoTA, with linear complexity.

📋 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

More detailed settings can be found in mamba3d.yaml.

🧾 2. Datasets

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

🥧 3. Training from scratch

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

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

To vote on ScanObjectNN/Modelnet40, run:

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

Few-shot learning, run:

bash run_fewshot.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 [email protected].

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