Home

Awesome

Stochastic Layer-Wise Shuffle: A Good Practice to Improve Vision Mamba Training

Official PyTorch implementation of ShuffleMamba, a regularization for improving Vision Mamba training.

<img src="Assets/framework.png" alt="framework" style="zoom: 67%;" />

Installation

For ImageNet1K classification training

# torch>=2.0, cuda>=11.8
pip install timm==0.4.12 mlflow==2.9.1
pip install causal-conv1d==1.1.0
pip install mamba-ssm==1.1.1

For ADE20K segmentation and COCO detection with Openmmlab tools

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

Training

Example for training the ShuffleMamba-B for 300 epochs: run the script run.sh. Example for multi-stage training for ShuffleMamba-Reg-B: run the script run_reg.sh. *all the models are training with 8 GPUs by default.

TO DO

Citation

@article{shufflemamba,
      title={Stochastic Layer-Wise Shuffle: A Good Practice to Improve Vision Mamba Training}, 
      author={Zizheng Huang and Haoxing Chen and Jiaqi Li and Jun Lan and Huijia Zhu and Weiqiang Wang and Limin Wang},
      journal={arXiv preprint arXiv:2408.17081},
      year={2024},
}

Acknowledgement

This repo is built based on Mamba-Reg, VideoMamba, VMamba and Vit-Adapter, thanks!