DOCSEARCHbyAlgolia
Submit Docs

MambaVision

repository·main·Indexed 25 days ago

https://github.com/nvlabs/mambavision

A hybrid vision backbone combining Mamba and Transformer architectures for high-performance computer vision tasks, including image classification, object detection, and semantic segmentation. The repository provides implementations for various backbones (Tiny, Small, Base) and integrates with the MMDetection framework for tasks like Cascade Mask R-CNN on the COCO dataset.

Tokens
8.5K
Snippets
17
Records
35
Agent score
81%

What's inside MambaVision

  1. Use MambaVision for Object Detection and Semantic Segmentation

    main

    MambaVision supports downstream tasks. The repository provides specific code and pretrained models for:

    • Object Detection: Uses Cascade Mask R-CNN. Available backbones include MambaVision-T-1K, MambaVision-S-1K, and MambaVision-B-1K.
    • Semantic Segmentation: Uses UPerNet. Available backbones include MambaVision-T-1K, MambaVision-S-1K, MambaVision-B-1K, and MambaVision-L3-512-21K.

  3. Train MambaVision models

    main

    You can train MambaVision models using either multi-GPU (via Slurm) or single-GPU setups.

    Multi-GPU Training (8 GPUs)

    Use srun to launch training on 8 GPUs. You must provide a configuration file from the configs/mamba_vision/ directory.

    srun --gres=gpu:8 python tools/train.py configs/mamba_vision/<CONFIG_FILE>.py

    Note: For paper-scale experiments, larger models (except 'small') were trained using 2 nodes with 16 GPUs. You can use the provided Slurm-based training scripts in tools/ for these configurations.

    Single-GPU Training

    For debugging or quick tests, use the following command:

    env CUDA_VISIBLE_DEVICES=0 python tools/train.py ${CONFIG}
    # multi‑GPU training (8 GPUs)
srun --gres=gpu:8 python tools/train.py configs/mamba_vision/<CONFIG_FILE>.py

# single‑GPU training
env CUDA_VISIBLE_DEVICES=0 python tools/train.py ${CONFIG}

  5. Configure Swin Transformer with one-stage detectors like RetinaNet

    main

    When combining a Swin Transformer backbone with a one-stage detector (e.g., RetinaNet), you must account for the Layer Normalization present at the output of the Swin Transformer.

    To ensure correct feature pyramid construction, you must:

    1. Set start_level to 0 in the FPN (Feature Pyramid Network) configuration.
    2. Set the out_indices of the backbone to [1, 2, 3].

    Refer to the retinanet_swin-t-p4-w7_fpn_1x_coco.py configuration file for a working implementation pattern.

  6. Install MambaVision for Semantic Segmentation

    main

    MambaVision semantic segmentation is built on top of the MMSegmentation framework.

    Prerequisites

    Ensure your environment is configured to run MambaVision pre-trained backbones by following the main MambaVision installation guide.

    Recommended Environment

    ComponentVersion
    PyTorch2.4.1+cu124
    CUDA12.4
    OpenCV4.10.0
    MMCV2.1.0
    MMSegmentation1.2.2
    MMEngine0.10.1

    Dependency Installation

    Install the required packages using pip:

    pip install \
  mmengine==0.10.1 \
  mmcv==2.1.0 \
  opencv-python-headless \
  mmsegmentation==1.2.2 \
  mmdet==3.3.0 \
  mmpretrain==1.2.0

  7. Install MambaVision for Object Detection

    main

    MambaVision object detection is built on the MMDetection framework.

    Prerequisites

    1. Verify Backbone Support: Ensure your environment is configured for MambaVision pre-trained backbones by following the main MambaVision installation guide.
    2. Recommended Environment:
      • PyTorch: 2.4.1+cu124
      • CUDA: 12.4
      • OpenCV: 4.10.0
      • MMCV: 2.1.0
      • MMDetection: 3.3.0
      • MMEngine: 0.10.1

    Dependency Installation

    Install the required packages using pip:

    pip install \
  mmengine==0.10.1 \
  mmcv==2.1.0 \
  opencv-python-headless \
  mmdet==3.3.0 \
  mmsegmentation==1.2.2 \
  mmpretrain==1.2.0

  9. Convert official Swin Transformer models to MMSegmentation style

    main

    If you are using pretrained models from the official Swin Transformer repository, you must convert their keys to match the MMSegmentation style. The project provides a conversion script swin2mmseg.py located in the tools/model_converters/ directory.

    To convert a model, run the script with the path to the original pretrained checkpoint as the first argument and the desired destination path as the second argument.

    python tools/model_converters/swin2mmseg.py ${PRETRAIN_PATH} ${STORE_PATH}

  10. Prepare ADE20K dataset for MambaVision

    main

    To use the ADE20K dataset for semantic segmentation, follow these steps:

    1. Download and Unzip

    cd <path-to-mambavision_seg-root>
mkdir -p data/ade20k && cd data/ade20k
wget http://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip
unzip ADEChallengeData2016.zip

    2. Verify Directory Structure

    Ensure your data is organized as follows:

    data/ade20k/
├── images/
│   ├── training/
│   └── validation/
└── annotations/
    ├── training/
    └── validation/

    3. Configure MMSegmentation

    Update your configuration file (e.g., in configs/mamba_vision/segmentation) to point to the data root and define the dataset types:

    data_root = 'data/ade20k/'

data = dict(
    train=dict(
        type='ADE20KDataset',
        data_root=data_root,
        img_dir='images/training',
        ann_dir='annotations/training'),
    val=dict(
        type='ADE20KDataset',
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation'),
    test=dict(
        type='ADE20KDataset',
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation')
)

  11. Prepare COCO Dataset for MambaVision Detection

    main

    To use the COCO dataset for object detection and instance segmentation, follow these steps:

    1. Download and Unzip

    Navigate to your MMDetection root and download the required files:

    cd <path-to-mmdetection-root>
mkdir -p data/coco && cd data/coco
wget http://images.cocodataset.org/zips/train2017.zip
wget http://images.cocodataset.org/zips/val2017.zip
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip
unzip train2017.zip
unzip val2017.zip
unzip annotations_trainval2017.zip

    2. Verify Directory Structure

    Ensure your data/coco/ directory looks like this:

    • annotations/: contains instances_train2017.json, instances_val2017.json, and optionally panoptic_train2017.json.
    • train2017/
    • val2017/

    3. Configure MMDetection Config

    Update your configuration file to point to the correct data_root and file paths:

    data_root = 'data/coco/'

data = dict(
    train=dict(
        img_prefix=data_root + 'train2017/',
        ann_file=data_root + 'annotations/instances_train2017.json'),
    val=dict(
        img_prefix=data_root + 'val2017/',
        ann_file=data_root + 'annotations/instances_val2017.json'),
    test=dict(
        img_prefix=data_root + 'val2017/',
        ann_file=data_root + 'annotations/instances_val2017.json')
)

  12. Evaluate MambaVision Object Detection Models

    main

    Evaluation can be performed using multi-GPU high-throughput inference or single-GPU debugging.

    Multi-GPU Inference (8 GPUs)

    Run the Slurm test script for high-throughput evaluation:

    bash tools/test.sh

    Single-GPU Inference

    Use the standard test.py tool. You must provide the config file, the checkpoint file, and specify the evaluation metrics (e.g., bbox and segm for instance segmentation):

    env CUDA_VISIBLE_DEVICES=0 \
  python tools/test.py \
    configs/mamba_vision/<CONFIG_FILE>.py \
    <CHECKPOINT_FILE>.pth \
    --eval bbox segm
    # multi‑GPU inference (8 GPUs)
bash tools/test.sh

# single‑GPU inference
env CUDA_VISIBLE_DEVICES=0 \
  python tools/test.py \
    configs/mamba_vision/<CONFIG_FILE>.py \
    <CHECKPOINT_FILE>.pth \
    --eval bbox segm