license: other
license_name: nvclv1
license_link: LICENSE
datasets:
- ILSVRC/imagenet-1k
pipeline_tag: image-classification
MambaVision: A Hybrid Mamba-Transformer Vision Backbone.
Model Overview
We introduce a novel mixer block by creating a symmetric path without SSM to enhance the modeling of global context. MambaVision has a hierarchical architecture that employs both self-attention and mixer blocks.
Model Performance
MambaVision demonstrates a strong performance by achieving a new SOTA Pareto-front in terms of Top-1 accuracy and throughput.
Model Usage
You must first login into HuggingFace to pull the model:
huggingface-cli login
It is highly recommended to install the requirements for MambaVision by running the following:
pip install mambavision
For each model, we offer two variants for image classification and feature extraction that can be imported with 1 line of code.
The model can be simply imported according to:
from transformers import AutoModelForImageClassification
model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
The model outputs logits when an image is passed. If label is additionally provided, cross entropy loss between the output prediction and label is computed.
The following demonstrates a minimal example of how to use the model:
from transformers import AutoModelForImageClassification
from PIL import Image
import requests
import torch
import timm
# import mambavision model
model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
# eval mode for inference
model.eval()
# prepare image for the model
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
# define a transform
transforms = timm.data.create_transform((3, 224, 224))
image = transforms(image).unsqueeze(0)
# put both model and image on cuda
model = model.cuda()
image = image.cuda()
# forward pass
outputs = model(image)
# You can then extract the predicted probabilities by applying softmax:
probabilities = torch.nn.functional.softmax(outputs['logits'], dim=0)
# In order to find the top 5 predicted class indexes and their corresponding values:
values, indices = torch.topk(probabilities, 5)