Memba for Vision Tasks

This directory includes the code for reproducing the vision task results in our paper.

Setup

Install dependencies. Our code is tested on CUDA 11.7, PyTorch 2.0.0, and Python 3.8.

pip install torch==2.0.0 torchvision==0.15.1
pip install -U openmim && mim install "mmpretrain==1.2"
pip install optuna einops spikingjelly

Install the modified Mamba package (from MambaPEFT):

pip install personal_lib/external_packages/mamba-1p1p1

Datasets and Weights

Download the VTab-1K dataset:

# Change BASEDIR in the script to your data directory
bash tools/download_vtab.sh

Download the Vim pre-trained weight from huggingface.co/hustvl/Vim-small-midclstok and place it as:

work_dirs/weights/backbone/vim/vim_s_midclstok_80p5acc.pth

Fine-tuning and Evaluation

Memba (Vim-S, LoRA on in+out proj + LIM)

CONFIG_NAME=mmpretrain/vim/vtab1k/1_small/memba_2_small_lora_in_out_proj_329616
CHECKPOINT_DIR=`date "+%Y%m%d_%H%M%S_%N"`
for DATASET in cifar caltech101 dtd oxford_flowers102 oxford_iiit_pet svhn sun397 \
               patch_camelyon eurosat resisc45 diabetic_retinopathy clevr_count clevr_dist \
               dmlab kitti dsprites_loc dsprites_ori smallnorb_azi smallnorb_ele
    do
        CUDA_VISIBLE_DEVICES=0 ./tools/dist_train.sh mmpretrain configs/${CONFIG_NAME}.py 1 \
            --work-dir work_dirs/${CONFIG_NAME}/${CHECKPOINT_DIR} \
            --cfg-options sub_dataset_name=${DATASET}
        CUDA_VISIBLE_DEVICES=0 ./tools/dist_test.sh mmpretrain configs/${CONFIG_NAME}.py \
            work_dirs/${CONFIG_NAME}/${CHECKPOINT_DIR}/last_checkpoint 1 \
            --cfg-options sub_dataset_name=${DATASET} | \
            tee work_dirs/${CONFIG_NAME}/${CHECKPOINT_DIR}/result_${DATASET}.txt
    done

You can override data_root at runtime if your data is in a different location:

--cfg-options sub_dataset_name=${DATASET} data_root=/your/data/path/

Available Configs

Configs for different Memba variants are in configs/mmpretrain/vim/vtab1k/:

• 1_small/ — Vim-Small configs
• 2_base/ — Vim-Base configs

Acknowledgements

Our code is based on Vim, MMPreTrain, MambaPEFT, and PETL-ViT.