EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications
EdgeNeXt effectively combines the strengths of both CNN and Transformer models and is a new efficient hybrid architecture. EdgeNeXt introduces a 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.[1]
Figure 1. Architecture of EdgeNeXt [1]
mindspore | ascend driver | firmware | cann toolkit/kernel |
---|---|---|---|
2.3.1 | 24.1.RC2 | 7.3.0.1.231 | 8.0.RC2.beta1 |
Please refer to the installation instruction in MindCV.
Please download the ImageNet-1K dataset for model training and validation.
- Distributed Training
It is easy to reproduce the reported results with the pre-defined training recipe. For distributed training on multiple Ascend 910 devices, please run
# distributed training on multiple NPU devices
msrun --bind_core=True --worker_num 8 python train.py --config configs/edgenext/edgenext_small_ascend.yaml --data_dir /path/to/imagenet
For detailed illustration of all hyper-parameters, please refer to config.py.
Note: As the global batch size (batch_size x num_devices) is an important hyper-parameter, it is recommended to keep the global batch size unchanged for reproduction or adjust the learning rate linearly to a new global batch size.
- Standalone Training
If you want to train or finetune the model on a smaller dataset without distributed training, please run:
# standalone training on single NPU device
python train.py --config configs/edgenext/edgenext_small_ascend.yaml --data_dir /path/to/dataset --distribute False
To validate the accuracy of the trained model, you can use validate.py
and parse the checkpoint path with --ckpt_path
.
python validate.py -c configs/edgenext/edgenext_small_ascend.yaml --data_dir /path/to/imagenet --ckpt_path /path/to/ckpt
Our reproduced model performance on ImageNet-1K is reported as follows.
Experiments are tested on ascend 910* with mindspore 2.3.1 graph mode.
model name | params(M) | cards | batch size | resolution | jit level | graph compile | ms/step | img/s | acc@top1 | acc@top5 | recipe | weight |
---|---|---|---|---|---|---|---|---|---|---|---|---|
edgenext_xx_small | 1.33 | 8 | 256 | 256x256 | O2 | 389s | 239.38 | 8555.43 | 70.64 | 89.75 | yaml | weights |
Experiments are tested on ascend 910 with mindspore 2.3.1 graph mode.
model name | params(M) | cards | batch size | resolution | jit level | graph compile | ms/step | img/s | acc@top1 | acc@top5 | recipe | weight |
---|---|---|---|---|---|---|---|---|---|---|---|---|
edgenext_xx_small | 1.33 | 8 | 256 | 256x256 | O2 | 311s | 191.24 | 10709.06 | 71.02 | 89.99 | yaml | weights |
- top-1 and top-5: Accuracy reported on the validation set of ImageNet-1K.
[1] Maaz M, Shaker A, Cholakkal H, et al. EdgeNeXt: efficiently amalgamated CNN-transformer architecture for Mobile vision applications[J]. arXiv preprint arXiv:2206.10589, 2022.