VMamba: Visual State Space Model

VMamba is a bidirectional state-space model finetuned on Imagenet dataset. It was introduced in the paper: VMamba: Visual State Space Model and was first released in this repo.

Disclaimer: This is not the official implementation, please refer to the official repo.

How to Get Started with the Model

Use the code below to get started with the model.

import torch
from PIL import Image
import torchvision.transforms as T
from transformers import AutoConfig, AutoModelForImageClassification

config = AutoConfig.from_pretrained('saurabhati/VMamba_ImageNet_82.6',trust_remote_code=True)
vmamba_model = AutoModelForImageClassification.from_pretrained('saurabhati/VMamba_ImageNet_82.6',trust_remote_code=True)

preprocess = T.Compose([
    T.Resize(224, interpolation=Image.BICUBIC),
    T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize(
        mean=[0.4850, 0.4560, 0.4060],
        std=[0.2290, 0.2240, 0.2250]
    )])

input_image = Image.open('/data/sls/scratch/sbhati/data/Imagenet/train/n02009912/n02009912_16160.JPEG')
input_image = preprocess(input_image)

with torch.no_grad():
    logits = vmamba_model(input_image.unsqueeze(0)).logits

predicted_label = vmamba_model.config.id2label[logits.argmax().item()]
predicted_label
'crane'

Citation

@article{liu2024vmamba,
  title={VMamba: Visual State Space Model},
  author={Liu, Yue and Tian, Yunjie and Zhao, Yuzhong and Yu, Hongtian and Xie, Lingxi and Wang, Yaowei and Ye, Qixiang and Liu, Yunfan},
  journal={arXiv preprint arXiv:2401.10166},
  year={2024}
}