Home

Awesome

EdgeNeXt

EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications [CADL'22, ECCVW]

Muhammad Maaz, Abdelrahman Shaker, Hisham Cholakkal, Salman Khan, Syed Waqas Zamir, Rao Muhammad Anwer and Fahad Shahbaz Khan

Website paper video slides

:rocket: News

<hr />

main figure

Abstract: In the pursuit of achieving ever-increasing accuracy, large and complex neural networks are usually developed. Such models demand high computational resources and therefore cannot be deployed on edge devices. It is of great interest to build resource-efficient general purpose networks due to their usefulness in several application areas. In this work, we strive to effectively combine the strengths of both CNN and Transformer models and propose a new efficient hybrid architecture EdgeNeXt. Specifically in EdgeNeXt, we introduce split depth-wise transpose attention (SDTA) encoder that splits input tensors into multiple channel groups and utilizes depth-wise convolution along with self-attention across channel dimensions to implicitly increase the receptive field and encode multi-scale features. Our extensive experiments on classification, detection and segmentation tasks, reveal the merits of the proposed approach, outperforming state-of-the-art methods with comparatively lower compute requirements. Our EdgeNeXt model with 1.3M parameters achieves 71.2% top-1 accuracy on ImageNet-1K, outperforming MobileViT with an absolute gain of 2.2% with 28% reduction in FLOPs. Further, our EdgeNeXt model with 5.6M parameters achieves 79.4% top-1 accuracy on ImageNet-1K.

<hr />

Model Zoo

NameAcc@1#ParamsMAddsModel
edgenext_base_usi83.6818.51M3.84Gmodel
edgenext_base_IN21K83.3118.51M3.84Gmodel
edgenext_base82.4718.51M3.84Gmodel
edgenext_small_usi81.075.59M1.26Gmodel
edgenext_small79.415.59M1.26Gmodel
edgenext_x_small74.962.34M538Mmodel
edgenext_xx_small71.231.33M261Mmodel
edgenext_small_bn_hs78.395.58M1.25Gmodel
edgenext_x_small_bn_hs74.872.34M536Mmodel
edgenext_xx_small_bn_hs70.331.33M260Mmodel
<hr />

Comparison with SOTA ViTs and Hybrid Architectures

results

<hr />

Comparison with Previous SOTA MobileViT (ICLR-2022)

results

<hr />

Qualitative Results (Segmentation)

results

Installation

  1. Create conda environment
conda create --name edgenext python=3.8
conda activate edgenext
  1. Install PyTorch and torchvision
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
  1. Install other dependencies
pip install -r requirements.txt
<hr />

Dataset Preparation

Download the ImageNet-1K classification dataset and structure the data as follows:

/path/to/imagenet-1k/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class2/
      img4.jpeg
<hr />

Evaluation

Download the pretrained weights and run the following command for evaluation on ImageNet-1K dataset.

wget https://github.com/mmaaz60/EdgeNeXt/releases/download/v1.0/edgenext_small.pth
python main.py --model edgenext_small --eval True --batch_size 16 --data_path <path to imagenet> --output_dir <results> --resume edgenext_small.pth

This should give,

Acc@1 79.412 Acc@5 94.512 loss 0.881
Note: For evaluating the USI model, please set usi_eval True.
<hr />

Training

On a single machine with 8 GPUs, run the following command to train EdgeNeXt-S model.

python -m torch.distributed.launch --nproc_per_node=8 main.py \
--model edgenext_small --drop_path 0.1 \
--batch_size 256 --lr 6e-3 --update_freq 2 \
--model_ema true --model_ema_eval true \
--data_path </path/to/imagenet-1k> \
--output_dir </path/to/save_results> \
--use_amp True --multi_scale_sampler
<hr />

Citation

If you use our work, please consider citing:

    @inproceedings{Maaz2022EdgeNeXt,
      title={EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications},
        author={Muhammad Maaz and Abdelrahman Shaker and Hisham Cholakkal and Salman Khan and Syed Waqas Zamir and Rao Muhammad Anwer and Fahad Shahbaz Khan},
      booktitle={International Workshop on Computational Aspects of Deep Learning at 17th European Conference on Computer Vision (CADL2022)},
      year={2022},
      organization={Springer}
}
<hr />

Contact

Should you have any question, please create an issue on this repository or contact at muhammad.maaz@mbzuai.ac.ae & abdelrahman.youssief@mbzuai.ac.ae

<hr />

References

Our code is based on ConvNeXt repository. We thank them for releasing their code.

Our Related Works