Awesome

SwiftFormer

SwiftFormer: Efficient Additive Attention for Transformer-based Real-time Mobile Vision Applications

Abdelrahman Shaker*1, Muhammad Maaz1, Hanoona Rasheed1, Salman Khan1, Ming-Hsuan Yang2,3 and Fahad Shahbaz Khan1,4

Mohamed Bin Zayed University of Artificial Intelligence1, University of California Merced2, Google Research3, Linkoping University4

:rocket: News

(Jul 14, 2023): SwiftFormer has been accepted at ICCV 2023. :fire::fire:
(Mar 27, 2023): Classification training and evaluation codes along with pre-trained models are released.

<hr /> <img src="images/Swiftformer_performance.png" width=60%> Comparison of our SwiftFormer Models with state-of-the-art on ImgeNet-1K. The latency is measured on iPhone 14 Neural Engine (iOS 16). <img src="images/attentions_comparison.png" width=99%> Comparison with different self-attention modules. (a) is a typical self-attention. (b) is the transpose self-attention, where the self-attention operation is applied across channel feature dimensions (d×d) instead of the spatial dimension (n×n). (c) is the separable self-attention of MobileViT-v2, it uses element-wise operations to compute the context vector from the interactions of Q and K matrices. Then, the context vector is multiplied by V matrix to produce the final output. (d) Our proposed efficient additive self-attention. Here, the query matrix is multiplied by learnable weights and pooled to produce global queries. Then, the matrix K is element-wise multiplied by the broadcasted global queries, resulting the global context representation. <details> <summary> Abstract </summary> Self-attention has become a defacto choice for capturing global context in various vision applications. However, its quadratic computational complexity with respect to image resolution limits its use in real-time applications, especially for deployment on resource-constrained mobile devices. Although hybrid approaches have been proposed to combine the advantages of convolutions and self-attention for a better speed-accuracy trade-off, the expensive matrix multiplication operations in self-attention remain a bottleneck. In this work, we introduce a novel efficient additive attention mechanism that effectively replaces the quadratic matrix multiplication operations with linear element-wise multiplications. Our design shows that the key-value interaction can be replaced with a linear layer without sacrificing any accuracy. Unlike previous state-of-the-art methods, our efficient formulation of self-attention enables its usage at all stages of the network. Using our proposed efficient additive attention, we build a series of models called "SwiftFormer" which achieves state-of-the-art performance in terms of both accuracy and mobile inference speed. Our small variant achieves 78.5% top-1 ImageNet-1K accuracy with only 0.8~ms latency on iPhone 14, which is more accurate and 2x faster compared to MobileViT-v2. </details>

Classification on ImageNet-1K

Models

Model	Top-1 accuracy	#params	GMACs	Latency	Ckpt	CoreML
SwiftFormer-XS	75.7%	3.5M	0.6G	0.7ms	XS	XS
SwiftFormer-S	78.5%	6.1M	1.0G	0.8ms	S	S
SwiftFormer-L1	80.9%	12.1M	1.6G	1.1ms	L1	L1
SwiftFormer-L3	83.0%	28.5M	4.0G	1.9ms	L3	L3

Detection and Segmentation Qualitative Results

Latency Measurement

The latency reported in SwiftFormer for iPhone 14 (iOS 16) uses the benchmark tool from XCode 14.

SwiftFormer meets Android

Community-driven results with Samsung Galaxy S23 Ultra, with Qualcomm Snapdragon 8 Gen 2:

Export & profiler results of SwiftFormer_L1:

QNN 2.16 2.17 2.18
Latency (msec) 2.63 2.26 2.43
Export & profiler results of SwiftFormerEncoder block:

QNN 2.16 2.17 2.18
Latency (msec) 2.17 1.69 1.7

Refer to script above for details of the input & block parameters.

QNN	2.16	2.17	2.18
Latency (msec)	2.63	2.26	2.43

QNN	2.16	2.17	2.18
Latency (msec)	2.17	1.69	1.7

❓ Interested in reproducing the results above?

Refer to Issue #14 for details about exporting & profiling.

ImageNet

Prerequisites

conda virtual environment is recommended.

conda create --name=swiftformer python=3.9
conda activate swiftformer

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install timm
pip install coremltools==5.2.0

Data preparation

Download and extract ImageNet train and val images from http://image-net.org. The training and validation data are expected to be in the train folder and val folder respectively:

|-- /path/to/imagenet/
    |-- train
    |-- val

Single machine multi-GPU training

We provide training script for all models in dist_train.sh using PyTorch distributed data parallel (DDP).

To train SwiftFormer models on an 8-GPU machine:

sh dist_train.sh /path/to/imagenet 8

Note: specify which model command you want to run in the script. To reproduce the results of the paper, use 16-GPU machine with batch-size of 128 or 8-GPU machine with batch size of 256. Auto Augmentation, CutMix, MixUp are disabled for SwiftFormer-XS, and CutMix, MixUp are disabled for SwiftFormer-S.

Multi-node training

On a Slurm-managed cluster, multi-node training can be launched as

sbatch slurm_train.sh /path/to/imagenet SwiftFormer_XS

Note: specify slurm specific paramters in slurm_train.sh script.

Testing

We provide an example test script dist_test.sh using PyTorch distributed data parallel (DDP). For example, to test SwiftFormer-XS on an 8-GPU machine:

sh dist_test.sh SwiftFormer_XS 8 weights/SwiftFormer_XS_ckpt.pth

Citation

if you use our work, please consider citing us:

@InProceedings{Shaker_2023_ICCV,
    author    = {Shaker, Abdelrahman and Maaz, Muhammad and Rasheed, Hanoona and Khan, Salman and Yang, Ming-Hsuan and Khan, Fahad Shahbaz},
    title     = {SwiftFormer: Efficient Additive Attention for Transformer-based Real-time Mobile Vision Applications},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    year      = {2023},
}

Contact:

If you have any questions, please create an issue on this repository or contact at abdelrahman.youssief@mbzuai.ac.ae.

Acknowledgement

Our code base is based on LeViT and EfficientFormer repositories. We thank authors for their open-source implementation.

Our Related Works

EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications, CADL'22, ECCV. Paper | Code.