Awesome

EMO

Official PyTorch implementation of "Rethinking Mobile Block for Efficient Attention-based Models, ICCV'23".

Abstract: This paper focuses on developing modern, efficient, lightweight models for dense predictions while trading off parameters, FLOPs, and performance. Inverted Residual Block (IRB) serves as the infrastructure for lightweight CNNs, but no counterpart has been recognized by attention-based studies. This work rethinks lightweight infrastructure from efficient IRB and effective components of Transformer from a unified perspective, extending CNN-based IRB to attention-based models and abstracting a one-residual Meta Mobile Block (MMB) for lightweight model design. Following simple but effective design criterion, we deduce a modern Inverted Residual Mobile Block (iRMB) and build a ResNet-like Efficient MOdel (EMO) with only iRMB for down-stream tasks. Extensive experiments on ImageNet-1K, COCO2017, and ADE20K benchmarks demonstrate the superiority of our EMO over state-of-the-art methods, e.g., EMO-1M/2M/5M achieve 71.5, 75.1, and 78.4 Top-1 that surpass equal-order CNN-/Attention-based models, while trading-off the parameter, efficiency, and accuracy well: running 2.8-4.0×↑ faster than EdgeNeXt on iPhone14. Code is available.

Top: Abstracted unified Meta-Mobile Block from Multi-Head Self-Attention and Feed-Forward Network in Transformer as well as efficient Inverted Residual Block in MobileNet-v2. Absorbing the experience of light-weight CNN and Transformer, an efficient but effective EMO is designed based on deduced iRMB.<br> Bottom: Performance vs. FLOPs comparisons with SoTA Transformer-based methods.

Main results

Image Classification for ImageNet-1K:

Object Detection Performance Based on SSDLite for COCO2017:

Object Detection Performance Based on RetinaNet for COCO2017:

Semantic Segmentation Based on DeepLabv3 for ADE20k:

Semantic Segmentation Based on PSPNet for ADE20k:

Classification

Environments

conda install -y pytorch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 cudatoolkit=11.3 -c pytorch -c conda-forge
pip3 install timm==0.6.5 tensorboardX einops torchprofile fvcore
git clone https://github.com/NVIDIA/apex && cd apex && pip3 install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ (optional)

Prepare ImageNet-1K Dataset

Download and extract ImageNet-1K dataset in the following directory structure:

├── imagenet
    ├── train
        ├── n01440764
            ├── n01440764_10026.JPEG
            ├── ...
        ├── ...
    ├── train.txt (optional)
    ├── val
        ├── n01440764
            ├── ILSVRC2012_val_00000293.JPEG
            ├── ...
        ├── ...
    └── val.txt (optional)

Test

Test with 8 GPUs in one node:

python3 -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --use_env run.py -c configs/mobile/cls_emo -m test model.name=EMO_1M trainer.data.batch_size=2048 model.model_kwargs.checkpoint_path=resources/EMO-1M/net.pth

This should give Top-1: 71.498 (Top-5: 90.368)

python3 -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --use_env run.py -c configs/mobile/cls_emo -m test model.name=EMO_2M trainer.data.batch_size=2048 model.model_kwargs.checkpoint_path=resources/EMO-2M/net.pth

This should give Top-1 75.134 (Top-5: 92.184)

python3 -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --use_env run.py -c configs/mobile/cls_emo -m test model.name=EMO_5M trainer.data.batch_size=2048 model.model_kwargs.checkpoint_path=resources/EMO-5M/net.pth

This should give Top-1 78.422 (Top-5: 93.970)

python3 -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --use_env run.py -c configs/mobile/cls_emo -m test model.name=EMO_6M trainer.data.batch_size=2048 model.model_kwargs.checkpoint_path=resources/EMO-6M/net.pth

This should give Top-1 78.988 (Top-5: 94.162)

Train

Train with 8 GPUs in one node:

python3 -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --use_env run.py -c configs/mobile/cls_emo -m train model.name=EMO_1M trainer.data.batch_size=2048

Down-Stream Tasks

Object Detection

• Ref to MMDetection for the environments.
• Configs can be found in down-stream-tasks/mmdetection/configs/ssd_emo and down-stream-tasks/mmdetection/configs/retinanet_emo
• E.g.:<br> ./tools/dist_train.sh configs/ssd_emo/ssdlite_emo_5M_pretrain_coco.py $GPU_NUM # for SSDLite with EMO-5M<br> ./tools/dist_train.sh configs/retinanet_emo/retinanet_emo_5M_fpn_1x_coco.py $GPU_NUM # for RetinaNet with EMO-5M

Semantic Segmentation

• Ref to MMSegmentation for the environments.
• Configs can be found in down-stream-tasks/mmsegmentation/configs/deeplabv3_emo and down-stream-tasks/mmsegmentation/configs/pspnet_emo
• E.g.:<br> ./tools/dist_train.sh configs/deeplabv3_emo/deeplabv3_emo_5M_pretrain_512x512_80k_ade20k.py $GPU_NUM # for DeepLabv3 with EMO-5M<br> ./tools/dist_train.sh configs/pspnet_emo/pspnet_emo_5M_512x512_80k_ade20k.py $GPU_NUM # for PSPNet with EMO-5M

Mobile Evaluation

• run python3 model/_emo_ios.py to obtain Core ML models
• Ref to Classifying Images with Vision and Core ML to build IOS App and test.

Citation

If our work is helpful for your research, please consider citing:

@inproceedings{emo,
  title={Rethinking Mobile Block for Efficient Attention-based Models},
  author={Zhang, Jiangning and Li, Xiangtai and Li, Jian and Liu, Liang and Xue, Zhucun and Zhang, Boshen and Jiang, Zhengkai and Huang, Tianxin and Wang, Yabiao and Wang, Chengjie},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={1389--1400},
  year={2023}
}

Acknowledgements

We thank but not limited to following repositories for providing assistance for our research:

• TIMM
• MMDetection
• MMSegmentation