Home

Awesome

2DMamba: Efficient State Space Model for Image Representation with Applications on Giga-Pixel Whole Slide Image Classification

Pytorch implementation for the 2DMamba framework described in the paper 2DMamba: Efficient State Space Model for Image Representation with Applications on Giga-Pixel Whole Slide Image Classification, arxiv.

<div> <img src="misc/overview_github.jpg" width="100%" alt="The overview of our framework."/> </div>

Installation

Install Anaconda/miniconda.
Required packages:

  $ conda create --name 2dmamba python=3.10
  $ conda activate 2dmamba
  $ conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia
  $ conda install nvidia/label/cuda-11.8.0::cuda
  $ conda install cmake # If you already have cmake, ignore this
  $ pip install pandas opencv-contrib-python kornia gpustat albumentations triton timm==0.4.12 tqdm pytest chardet yacs termcolor submitit tensorboardX fvcore seaborn einops tensorboard joblib
  $	pip install mamba-ssm
  $ pip install --force-reinstall -v "numpy==1.26.4"
  $ pip install mmengine==0.10.1 mmcv==2.1.0 opencv-python-headless ftfy regex
  $ pip install mmdet==3.3.0 mmsegmentation==1.2.2 mmpretrain==1.2.0

You can also use docker or singularity. We provide the Dockerfile we used in our experiments. You can also pull our image on dockerhub, which should be identical to our environment, by:

  $ docker pull skykiny/mamba # For docker
  $ singularity pull docker://skykiny/mamba:latest # For singularity

Build 2DMamba CUDA kernel

We use CMake to build our CUDA kernel. Please replace the -DPython_ROOT_DIR="/opt/conda" in cuda_kernel/build.sh with your python root directory. E.g. if you use conda environment and your python is located at /home/jzhang/Dev/anaconda3_2023/envs/vmamba/bin/python, you should set -DPython_ROOT_DIR="/home/jzhang/Dev/anaconda3_2023/envs/vmamba". Then run bash build.sh, the compiled pscan.so should appear under v2dmamba_scan folder. You can try cd to the root directory of this project and run import v2dmamba_scan in python to verify if it is correct.

Contact

If you have any questions or concerns, feel free to report an issue or directly contact us at Jingwei Zhang jingwezhang@cs.stonybrook.edu, Anh Tien Nguyen tienanhnguyen9991@gmail.com or Xi Han xihan1@cs.stonybrook.edu.

Acknowledgments

Our framework is based on Mamba, VMamba and mamba.py. Thanks for their outstanding code.

Citation

If you use the code or results in your research, please use the following BibTeX entry.

@article{zhang20242dmamba,
  title={2DMamba: Efficient State Space Model for Image Representation with Applications on Giga-Pixel Whole Slide Image Classification},
  author={Zhang, Jingwei and Nguyen, Anh Tien and Han, Xi and Trinh, Vincent Quoc-Huy and Qin, Hong and Samaras, Dimitris and Hosseini, Mahdi S},
  journal={arXiv preprint arXiv:2412.00678},
  year={2024}
}