|
--- |
|
datasets: |
|
- imagenet-1k |
|
library_name: transformers |
|
pipeline_tag: image-classification |
|
license: other |
|
tags: |
|
- vision |
|
- image-classification |
|
--- |
|
|
|
# MobileViTv2 (mobilevitv2-1.0-imagenet1k-256) |
|
|
|
<!-- Provide a quick summary of what the model is/does. --> |
|
MobileViTv2 is the second version of MobileViT. It was proposed in [Separable Self-attention for Mobile Vision Transformers](https://arxiv.org/abs/2206.02680) by Sachin Mehta and Mohammad Rastegari, and first released in [this](https://github.com/apple/ml-cvnets) repository. The license used is [Apple sample code license](https://github.com/apple/ml-cvnets/blob/main/LICENSE). |
|
|
|
Disclaimer: The team releasing MobileViT did not write a model card for this model so this model card has been written by the Hugging Face team. |
|
|
|
### Model Description |
|
|
|
<!-- Provide a longer summary of what this model is. --> |
|
MobileViTv2 is constructed by replacing the multi-headed self-attention in MobileViT with separable self-attention. |
|
|
|
### Intended uses & limitations |
|
|
|
You can use the raw model for image classification. See the [model hub](https://huggingface.co/models?search=mobilevitv2) to look for fine-tuned versions on a task that interests you. |
|
|
|
### How to use |
|
|
|
Here is how to use this model to classify an image of the COCO 2017 dataset into one of the 1,000 ImageNet classes: |
|
|
|
```python |
|
from transformers import MobileViTv2FeatureExtractor, MobileViTv2ForImageClassification |
|
from PIL import Image |
|
import requests |
|
|
|
url = "http://images.cocodataset.org/val2017/000000039769.jpg" |
|
image = Image.open(requests.get(url, stream=True).raw) |
|
|
|
feature_extractor = MobileViTv2FeatureExtractor.from_pretrained("shehan97/mobilevitv2-1.0-imagenet1k-256") |
|
model = MobileViTv2ForImageClassification.from_pretrained("shehan97/mobilevitv2-1.0-imagenet1k-256") |
|
|
|
inputs = feature_extractor(images=image, return_tensors="pt") |
|
|
|
outputs = model(**inputs) |
|
logits = outputs.logits |
|
|
|
# model predicts one of the 1000 ImageNet classes |
|
predicted_class_idx = logits.argmax(-1).item() |
|
print("Predicted class:", model.config.id2label[predicted_class_idx]) |
|
``` |
|
|
|
Currently, both the feature extractor and model support PyTorch. |
|
|
|
## Training data |
|
|
|
The MobileViT model was pretrained on [ImageNet-1k](https://huggingface.co/datasets/imagenet-1k), a dataset consisting of 1 million images and 1,000 classes. |
|
|
|
|
|
### BibTeX entry and citation info |
|
|
|
```bibtex |
|
@inproceedings{vision-transformer, |
|
title = {Separable Self-attention for Mobile Vision Transformers}, |
|
author = {Sachin Mehta and Mohammad Rastegari}, |
|
year = {2022}, |
|
URL = {https://arxiv.org/abs/2206.02680} |
|
} |
|
``` |
