/efficientvit

EfficientViT is a new family of vision models for efficient high-resolution vision.

Primary LanguagePythonApache License 2.0Apache-2.0

EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction (paper, poster)

News

If you are interested in getting updates, please join our mailing list here.

  • [2023/11/20] EfficientViT is available in the NVIDIA Jetson Generative AI Lab.
  • [2023/11/20] We will soon release the second version of EfficientViT SAM models.
  • [2023/10/22] ImageNet training scripts for the EfficientViT L series have been released.
  • [2023/09/18] EfficientViT for Segment Anything Model (SAM) is released. EfficientViT SAM runs at 1009 images/s on A100 GPU, compared to ViT-H (12 images/s), mobileSAM (297 images/s), and nanoSAM (744 image/s, but much lower mIoU)
  • [2023/09/12] EfficientViT is highlighted by MIT home page and MIT News.
  • [2023/07/18] EfficientViT is accepted by ICCV 2023.

Demo

EfficientViT-L0 for Segment Anything (1009 image/s on A100 GPU) demo demo demo

EfficientViT-L1 for Semantic Segmentation (45.9ms on Nvidia Jetson AGX Orin, 82.716 mIoU on Cityscapes)

demo

About EfficientViT Models

EfficientViT is a new family of ViT models for efficient high-resolution dense prediction vision tasks. The core building block of EfficientViT is a lightweight, multi-scale linear attention module that achieves global receptive field and multi-scale learning with only hardware-efficient operations, making EfficientViT TensorRT-friendly and suitable for GPU deployment.

Third-Party Implementation/Integration

Getting Started

Installation

conda create -n efficientvit python=3.10
conda activate efficientvit
conda install -c conda-forge mpi4py openmpi
pip install -r requirements.txt

Dataset

ImageNet: https://www.image-net.org/
Our code expects the ImageNet dataset directory to follow the following structure:

imagenet
├── train
├── val
Cityscapes: https://www.cityscapes-dataset.com/
Our code expects the Cityscapes dataset directory to follow the following structure:

cityscapes
├── gtFine
|   ├── train
|   ├── val
├── leftImg8bit
|   ├── train
|   ├── val
ADE20K: https://groups.csail.mit.edu/vision/datasets/ADE20K/
Our code expects the ADE20K dataset directory to follow the following structure:

ade20k
├── annotations
|   ├── training
|   ├── validation
├── images
|   ├── training
|   ├── validation

Pretrained Models

Latency/Throughput is measured on NVIDIA Jetson Nano, NVIDIA Jetson AGX Orin, and NVIDIA A100 GPU with TensorRT, fp16. Data transfer time is included.

Segment Anything

In this version, the EfficientViT segment anything models are trained using the image embedding extracted by SAM ViT-H as the target. The prompt encoder and mask decoder are the same as SAM ViT-H.

Image Encoder COCO-val2017 mIoU (all) COCO-val2017 mIoU (large) COCO-val2017 mIoU (medium) COCO-val2017 mIoU (small) Params MACs A100 Throughput Checkpoint
NanoSAM 70.6 79.6 73.8 62.4 - - 744 image/s -
MobileSAM 72.8 80.4 75.9 65.8 - - 297 image/s -
EfficientViT-L0-SAM 74.454 81.410 77.201 68.159 31M 35G 1009 image/s link
EfficientViT-L1-SAM 75.183 81.786 78.110 68.944 44M 49G 815 image/s link
EfficientViT-L2-SAM 75.656 81.706 78.644 69.689 57M 69G 634 image/s link

ImageNet

All EfficientViT classification models are trained on ImageNet-1K with random initialization (300 epochs + 20 warmup epochs) using supervised learning.

Model Resolution ImageNet Top1 Acc ImageNet Top5 Acc Params MACs A100 Throughput Checkpoint
EfficientNetV2-S 384x384 83.9 - 22M 8.4G 2869 image/s -
EfficientNetV2-M 480x480 85.2 - 54M 25G 1160 image/s -
EfficientViT-L1 224x224 84.484 96.862 53M 5.3G 6207 image/s link
EfficientViT-L2 224x224 85.050 97.090 64M 6.9G 4998 image/s link
EfficientViT-L2 256x256 85.366 97.216 64M 9.1G 3969 image/s link
EfficientViT-L2 288x288 85.630 97.364 64M 11G 3102 image/s link
EfficientViT-L2 320x320 85.734 97.438 64M 14G 2525 image/s link
EfficientViT-L2 384x384 85.978 97.518 64M 20G 1784 image/s link
EfficientViT-L3 224x224 85.814 97.198 246M 28G 2081 image/s link
EfficientViT-L3 256x256 85.938 97.318 246M 36G 1641 image/s link
EfficientViT-L3 288x288 86.070 97.440 246M 46G 1276 image/s link
EfficientViT-L3 320x320 86.230 97.474 246M 56G 1049 image/s link
EfficientViT-L3 384x384 86.408 97.632 246M 81G 724 image/s link
EfficientViT B series
Model Resolution ImageNet Top1 Acc ImageNet Top5 Acc Params MACs Jetson Nano (bs1) Jetson Orin (bs1) Checkpoint
EfficientViT-B1 224x224 79.390 94.346 9.1M 0.52G 24.8ms 1.48ms link
EfficientViT-B1 256x256 79.918 94.704 9.1M 0.68G 28.5ms 1.57ms link
EfficientViT-B1 288x288 80.410 94.984 9.1M 0.86G 34.5ms 1.82ms link
EfficientViT-B2 224x224 82.100 95.782 24M 1.6G 50.6ms 2.63ms link
EfficientViT-B2 256x256 82.698 96.096 24M 2.1G 58.5ms 2.84ms link
EfficientViT-B2 288x288 83.086 96.302 24M 2.6G 69.9ms 3.30ms link
EfficientViT-B3 224x224 83.468 96.356 49M 4.0G 101ms 4.36ms link
EfficientViT-B3 256x256 83.806 96.514 49M 5.2G 120ms 4.74ms link
EfficientViT-B3 288x288 84.150 96.732 49M 6.5G 141ms 5.63ms link

Cityscapes

Model Resolution Cityscapes mIoU Params MACs Jetson Orin Latency (bs1) A100 Throughput (bs1) Checkpoint
EfficientViT-L1 1024x2048 82.716 40M 282G 45.9ms 122 image/s link
EfficientViT-L2 1024x2048 83.228 53M 396G 60.0ms 102 image/s link
EfficientViT B series
Model Resolution Cityscapes mIoU Params MACs Jetson Nano (bs1) Jetson Orin (bs1) Checkpoint
EfficientViT-B0 1024x2048 75.653 0.7M 4.4G 275ms 9.9ms link
EfficientViT-B1 1024x2048 80.547 4.8M 25G 819ms 24.3ms link
EfficientViT-B2 1024x2048 82.073 15M 74G 1676ms 46.5ms link
EfficientViT-B3 1024x2048 83.016 40M 179G 3192ms 81.8ms link

ADE20K

Model Resolution ADE20K mIoU Params MACs Jetson Orin Latency (bs1) A100 Throughput (bs16) Checkpoint
EfficientViT-L1 512x512 49.191 40M 36G 7.2ms 947 image/s link
EfficientViT-L2 512x512 50.702 51M 45G 9.0ms 758 image/s link
EfficientViT B series
Model Resolution ADE20K mIoU Params MACs Jetson Nano (bs1) Jetson Orin (bs1) Checkpoint
EfficientViT-B1 512x512 42.840 4.8M 3.1G 110ms 4.0ms link
EfficientViT-B2 512x512 45.941 15M 9.1G 212ms 7.3ms link
EfficientViT-B3 512x512 49.013 39M 22G 411ms 12.5ms link

Usage

# segment anything
from efficientvit.sam_model_zoo import create_sam_model

efficientvit_sam = create_sam_model(
  name="l2", weight_url="assets/checkpoints/sam/l2.pt",
)
efficientvit_sam = efficientvit_sam.cuda().eval()
from efficientvit.models.efficientvit.sam import EfficientViTSamPredictor

efficientvit_sam_predictor = EfficientViTSamPredictor(efficientvit_sam)
from efficientvit.models.efficientvit.sam import EfficientViTSamAutomaticMaskGenerator

efficientvit_mask_generator = EfficientViTSamAutomaticMaskGenerator(efficientvit_sam)
# classification
from efficientvit.cls_model_zoo import create_cls_model

model = create_cls_model(
  name="l3", weight_url="assets/checkpoints/cls/l3-r384.pt"
)
# semantic segmentation
from efficientvit.seg_model_zoo import create_seg_model

model = create_seg_model(
  name="l2", dataset="cityscapes", weight_url="assets/checkpoints/seg/cityscapes/l2.pt"
)

model = create_seg_model(
  name="l2", dataset="ade20k", weight_url="assets/checkpoints/seg/ade20k/l2.pt"
)

Evaluation

Please run eval_sam_coco.py, eval_cls_model.py or eval_seg_model.py to evaluate our models.

Examples: segment anything, classification, segmentation

Visualization

Please run demo_sam_model.py to visualize our segment anything models.

Example:

# segment everything
python demo_sam_model.py --model l1 --mode all

# prompt with points
python demo_sam_model.py --model l1 --mode point

# prompt with box
python demo_sam_model.py --model l1 --mode box --box "[150,70,630,400]"

Please run eval_seg_model.py to visualize the outputs of our semantic segmentation models.

Example:

python eval_seg_model.py --dataset cityscapes --crop_size 1024 --model b3 --save_path demo/cityscapes/b3/

Export TFLite

To generate TFLite files, please refer to tflite_export.py. It requires the TinyNN package.

pip install git+https://github.com/alibaba/TinyNeuralNetwork.git

Example:

python tflite_export.py --export_path model.tflite --task seg --dataset ade20k --model b3 --resolution 512 512

Export ONNX

To generate ONNX files, please refer to onnx_export.py.

To export ONNX files for EfficientViT SAM models, please refer to the scripts shared by CVHub.

Training

Please see TRAINING.md for detailed training instructions.

Contact

Han Cai: hancai@mit.edu

TODO

  • ImageNet Pretrained models
  • Segmentation Pretrained models
  • ImageNet training code
  • EfficientViT L series, designed for cloud
  • EfficientViT for segment anything
  • EfficientViT for image generation
  • EfficientViT for super-resolution
  • Segmentation training code

Citation

If EfficientViT is useful or relevant to your research, please kindly recognize our contributions by citing our paper:

@article{cai2022efficientvit,
  title={Efficientvit: Enhanced linear attention for high-resolution low-computation visual recognition},
  author={Cai, Han and Gan, Chuang and Han, Song},
  journal={arXiv preprint arXiv:2205.14756},
  year={2022}
}