MobileViT

MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer, arxiv

PaddlePaddle training/validation code and pretrained models for MobileViT.

The official apple implementation is here.

This implementation is developed by PaddleViT.

MobileViT Transformer Model Overview

Update

• Update (2021-12-30): Add multi scale sampler DDP and update mobilevit_s model weights.
• Update (2021-11-02): Pretrained model weights (mobilevit_s) is released.
• Update (2021-11-02): Pretrained model weights (mobilevit_xs) is released.
• Update (2021-10-29): Pretrained model weights (mobilevit_xxs) is released.
• Update (2021-10-20): Initial code is released.

Models Zoo

Model	Acc@1	Acc@5	#Params	FLOPs	Image Size	Crop_pct	Interpolation	Link
mobilevit_xxs	70.31	89.68	1.32M	0.44G	256	1.0	bicubic	google/baidu(axpc)
mobilevit_xs	74.47	92.02	2.33M	0.95G	256	1.0	bicubic	google/baidu(hfhm)
mobilevit_s	76.74	93.08	5.59M	1.88G	256	1.0	bicubic	google/baidu(34bg)
mobilevit_s $\dag$	77.83	93.83	5.59M	1.88G	256	1.0	bicubic	google/baidu(92ic)

The results are evaluated on ImageNet2012 validation set.

All models are trained from scratch using PaddleViT.

$\dag$ means model is trained from scratch using PaddleViT using multi scale sampler DDP.

Notebooks

We provide a few notebooks in aistudio to help you get started:

*(coming soon)*

Requirements

• Python>=3.6
• yaml>=0.2.5
• PaddlePaddle>=2.1.0
• yacs>=0.1.8

Data

ImageNet2012 dataset is used in the following folder structure:

│imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

Usage

To use the model with pretrained weights, download the .pdparam weight file and change related file paths in the following python scripts. The model config files are located in ./configs/.

For example, assume the downloaded weight file is stored in ./mobilevit_xxs.pdparams, to use the mobilevit_xxs model in python:

from config import get_config
from mobile_vit import build_mobile_vit as build_model
import paddle
# config files in ./configs/
config = get_config('./configs/mobilevit_xxs.yaml')
# build model
model = build_model(config)
# load pretrained weights
model_state_dict = paddle.load('./mobilevit_xxs.pdparams')
model.set_dict(model_state_dict)

Evaluation

To evaluate MobileViT model performance on ImageNet2012 with a single GPU, run the following script using command line:

sh run_eval.sh

or

CUDA_VISIBLE_DEVICES=0 \
python main_single_gpu.py \
    -cfg=./configs/mobilevit_xxs.yaml \
    -dataset=imagenet2012 \
    -batch_size=16 \
    -data_path=/path/to/dataset/imagenet/val \
    -eval \
    -pretrained=/path/to/pretrained/model/mobilevit_xxs  # .pdparams is NOT needed

Run evaluation using multi-GPUs:

sh run_eval_multi.sh

or

CUDA_VISIBLE_DEVICES=0,1,2,3 \
python main_multi_gpu.py \
    -cfg=./configs/mobilevit_xxs.yaml \
    -dataset=imagenet2012 \
    -batch_size=16 \
    -data_path=/path/to/dataset/imagenet/val \
    -eval \
    -pretrained=/path/to/pretrained/model/mobilevit_xxs  # .pdparams is NOT needed

Training

To train the MobileVit XXS model on ImageNet2012 with single GPU, run the following script using command line:

sh run_train.sh

or

CUDA_VISIBLE_DEVICES=0 \
python main_singel_gpu.py \
  -cfg=./configs/mobilevit_xxs.yaml \
  -dataset=imagenet2012 \
  -batch_size=32 \
  -data_path=/path/to/dataset/imagenet/train

Run training using multi-GPUs:

sh run_train_multi.sh

or

CUDA_VISIBLE_DEVICES=0,1,2,3 \
python main_multi_gpu.py \
    -cfg=./configs/mobilevit_xxs.yaml \
    -dataset=imagenet2012 \
    -batch_size=16 \
    -data_path=/path/to/dataset/imagenet/train

Visualization Attention Map

(coming soon)

Reference

@article{mehta2021mobilevit,
  title={MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer},
  author={Mehta, Sachin and Rastegari, Mohammad},
  journal={arXiv preprint arXiv:2110.02178},
  year={2021}
}