🚀 VMamba: 视觉状态空间模型

VMamba 是一个在 ImageNet 数据集上进行微调的双向状态空间模型。它在论文 VMamba: Visual State Space Model 中被提出，并首次在 此仓库 中发布。

声明：这不是官方实现，请参考 官方仓库。这是我（Saurabhchand Bhati）为图像、音频分类任务添加 VMamba 主干网络的正在进行的工作。

🚀 快速开始

使用以下代码开始使用该模型。

基础用法

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'

📄 许可证

本项目使用 MIT 许可证。

📚 引用

@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}
}