一背景

端側(cè)場(chǎng)景通常對(duì)模型推理的實(shí)時(shí)性要求較高，但大部分輕量級(jí)的Vision Transformer網(wǎng)絡(luò)仍無(wú)法在已有邊緣側(cè)/端側(cè)AI設(shè)備（CPU、NPU）上達(dá)到輕量級(jí)CNN（如MobileNet）相媲美的速度。為了實(shí)現(xiàn)對(duì)ViT模型的實(shí)時(shí)部署，微軟和港中文共同在CVPR2023提出論文《EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention》。

本文將簡(jiǎn)單介紹EfficientViT的基本架構(gòu)原理，同時(shí)手把手帶領(lǐng)感興趣的朋友基于該論文Github項(xiàng)目導(dǎo)出ONNX模型，并將其部署在優(yōu)秀的端側(cè)AI芯片AX650N上，希望能給行業(yè)內(nèi)對(duì)邊緣側(cè)/端側(cè)部署Transformer模型的愛(ài)好者提供新的思路。

二EfficientViT

論文從三個(gè)維度分析了ViT的速度瓶頸，包括多頭自注意力（MHSA）導(dǎo)致的大量訪存時(shí)間，注意力頭之間的計(jì)算冗余，以及低效的模型參數(shù)分配，進(jìn)而提出了一個(gè)高效ViT模型EfficientViT。它以EfficientViT block作為基礎(chǔ)模塊，每個(gè)block由三明治結(jié)構(gòu)（Sandwich Layout）和級(jí)聯(lián)組注意力（Cascaded Group Attention, CGA）組成。在論文中作者進(jìn)一步進(jìn)行了參數(shù)重分配（Parameter Reallocation）以實(shí)現(xiàn)更高效的Channel、Block和Stage數(shù)量權(quán)衡。該方法在ImageNet數(shù)據(jù)集上實(shí)現(xiàn)了77.1%的Top-1分類準(zhǔn)確率，超越了MobileNetV3-Large 1.9%精度的同時(shí)，在NVIDIA V100 GPU 和 Intel Xeon CPU上實(shí)現(xiàn)了40.4%和45.2%的吞吐量提升，并且大幅領(lǐng)先其他輕量級(jí)ViT的速度和精度。

● 論文：

https://arxiv.org/abs/2305.070271

● Github鏈接：

https://github.com/microsoft/Cream/tree/main/EfficientViT1

2.1 骨干網(wǎng)絡(luò)

該論文的核心貢獻(xiàn)是提出了一種EfficientViT block，每個(gè)EfficientViT block的輸入特征先經(jīng)過(guò)N個(gè)FFN，再經(jīng)過(guò)一個(gè)級(jí)聯(lián)組注意力CGA，再經(jīng)過(guò)N個(gè)FFN層變換得到輸出特征。這一基礎(chǔ)模塊減少了注意力的使用，緩解了注意力計(jì)算導(dǎo)致的訪存時(shí)間消耗問(wèn)題。同時(shí)，作者在每個(gè)FFN之前加入了一層DWConv作為局部token之間信息交互并幫助引入歸納偏置。

另外考慮到BN可以與FC和Conv在推理時(shí)融合以實(shí)現(xiàn)加速。因此使用Batch Normalization替換Layer Normalization，同時(shí)在大尺度下層數(shù)更少，并在每個(gè)stage用了小于2的寬度擴(kuò)展系數(shù)以減輕深層的冗余。

qkv-backbone

benchmark

三AX650N

AX650N是一款兼具高算力與高能效比的SoC芯片，集成了八核Cortex-A55 CPU，10.8TOPs@INT8 NPU，支持8K@30fps的ISP，以及H.264、H.265編解碼的VPU。接口方面，AX650N支持64bit LPDDR4x，多路MIPI輸入，千兆Ethernet、USB、以及HDMI 2.0b輸出，并支持32路1080p@30fps解碼。強(qiáng)大的性能可以讓AX650N幫助用戶在智慧城市，智慧教育，智能制造等領(lǐng)域發(fā)揮更大的價(jià)值。

四模型轉(zhuǎn)換

本文以EfficientViT-M5為例。

4.1 Pulsar2

Pulsar2是新一代AI工具鏈，包含模型轉(zhuǎn)換、離線量化、模型編譯、異構(gòu)調(diào)度四合一超強(qiáng)功能，進(jìn)一步強(qiáng)化了網(wǎng)絡(luò)模型高效部署的需求。在針對(duì)第三NPU架構(gòu)進(jìn)行了深度定制優(yōu)化的同時(shí)，也擴(kuò)展了算子&模型支持的能力及范圍，對(duì)Transformer結(jié)構(gòu)的網(wǎng)絡(luò)也有較好的支持。

npu-software-arch

4.2 模型下載

下載github倉(cāng)庫(kù)，并進(jìn)入到EfficientViT/Classification路徑下安裝相關(guān)依賴庫(kù)

下載安裝

$ git clone https://github.com/microsoft/Cream.git
$ cd Cream/EfficientViT/classification/
$ pip install -r requirements.txt

●獲取PyTorch模型

雖然官方倉(cāng)庫(kù)中已經(jīng)提供了ONNX模型，但是我們發(fā)現(xiàn)其ONNX模型的Batch Size=16，并不適合普通的端側(cè)芯片進(jìn)行評(píng)估，因此這里選擇從pth文件重新生成Batch Size=1的ONNX模型，更利于端側(cè)芯片部署。

●下載pth文件

下載pth文件

$ wget https://github.com/xinyuliu-jeffrey/EfficientViT_Model_Zoo/releases/download/v1.0/efficientvit_m5.pth

●導(dǎo)出ONNX模型并初步優(yōu)化計(jì)算圖

導(dǎo)出onnx

$ python export_onnx_efficientvit_m5.py
$ onnxsim efficientvit_m5.onnx efficientvit_m5-sim.onnx

● export_onnx_efficientvit_m5.py源碼如下所示。

export_onnx源碼

from model import build
from timm.models import create_model
import torch


model = create_model(
    "EfficientViT_M5",
    num_classes=1000,
    distillation=False,
    pretrained=False,
    fuse=False,
  )


checkpoint = torch.load("./efficientvit_m5.pth", map_location='cpu')
state_dict = checkpoint['model']
model.load_state_dict(state_dict)
model.eval()
dummy_input = torch.rand([1,3,224,224])


model(dummy_input)


torch.onnx.export(model, dummy_input, "efficientvit_m5.onnx", opset_version=11)

4.3 模型編譯

一鍵完成圖優(yōu)化、離線量化、編譯、對(duì)分功能。整個(gè)過(guò)程耗時(shí)只需20秒。

編譯log

$ pulsar2 build --input model/efficientvit_m5-sim.onnx --output_dir efficientvit-m5/ --config config/effientvit_config.json
Building onnx ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
patool: Extracting ./dataset/imagenet-32-images.tar ...
patool: running /usr/bin/tar --extract --file ./dataset/imagenet-32-images.tar --directory efficientvit-m5/quant/dataset/input_1
patool: ... ./dataset/imagenet-32-images.tar extracted to `efficientvit-m5/quant/dataset/input_1'.
                                           Quant Config Table
┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Input  ┃ Shape      ┃ Dataset Directory ┃ Data Format ┃ Tensor Format ┃ Mean                          ┃ Std                       ┃
┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ input.1 │ [1, 3, 224, 224] │ input_1      │ Image    │ RGB      │ [123.67500305175781, 116.27999877929688,        │ [58.39500045776367, 57.119998931884766, 57.375] │
│     │         │          │       │        │ 103.52999877929688]                   │                         │
└─────────┴──────────────────┴───────────────────┴─────────────┴───────────────┴─────────────────────────────────────────────────────────┴─────────────────────────────────────────────────┘
Transformer optimize level: 2
32 File(s) Loaded.
[14:22:37] AX LSTM Operation Format Pass Running ...   Finished.
[14:22:37] AX Refine Operation Config Pass Running ...  Finished.
[14:22:37] AX Transformer Optimize Pass Running ...    Finished.
[14:22:37] AX Reset Mul Config Pass Running ...      Finished.
[14:22:37] AX Tanh Operation Format Pass Running ...   Finished.
[14:22:37] AX Softmax Format Pass Running ...       Finished.
[14:22:37] AX Quantization Config Refine Pass Running ... Finished.
[14:22:37] AX Quantization Fusion Pass Running ...    Finished.
[14:22:37] AX Quantization Simplify Pass Running ...   Finished.
[14:22:37] AX Parameter Quantization Pass Running ...   Finished.
Calibration Progress(Phase 1): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:03<00:00, ?8.60it/s]
Finished.
[1442] AX Passive Parameter Quantization Running ... ?Finished.
[1442] AX Parameter Baking Pass Running ... ? ? ? ? ? Finished.
[1442] AX Refine Int Parameter pass Running ... ? ? ? Finished.
quant.axmodel export success: efficientvit-m5/quant/quant_axmodel.onnx
Building native ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 000
......
2023-05-19 1448.523 | INFO ? ? | yasched.test_onepass1438 - max_cycle = 1052639
2023-05-19 14:22:55.172 | INFO ? ? | yamain.command.build890 - fuse 1 subgraph(s)

4.4 Graph Optimize

這里主要是對(duì)輸出頭的BN+FC的結(jié)構(gòu)進(jìn)行了簡(jiǎn)單的融合，利于后續(xù)編譯階段提升執(zhí)行效率。

圖優(yōu)化-前

圖優(yōu)化+量化后

五上板部署

5.1 AX-Samples

開(kāi)源項(xiàng)目AX-Samples實(shí)現(xiàn)了常見(jiàn)的深度學(xué)習(xí)開(kāi)源算法在愛(ài)芯元智的AI SoC上的示例代碼，方便社區(qū)開(kāi)發(fā)者進(jìn)行快速評(píng)估和適配。最新版本已逐步完善基于AX650系列的NPU示例，其中Classification通用示例可直接運(yùn)行前面章節(jié)編譯生成的EfficientViT模型。

https://github.com/AXERA-TECH/ax-samples/blob/main/examples/ax650/ax_classification_steps.cc

5.2 運(yùn)行

運(yùn)行l(wèi)og

--------------------------------------
/opt/qtang # sample_npu_classification -m efficientvit-m5-npu1.axmodel -i cat.jpg -r 10
--------------------------------------
model file : efficientvit-m5-npu1.axmodel
image file : cat.jpg
img_h, img_w : 224 224
--------------------------------------
topk cost time:0.07 ms
5.5997, 285
5.3721, 283
5.0079, 281
4.5982, 284
4.1884, 282
--------------------------------------
Repeat 10 times, avg time 1.24 ms, max_time 1.24 ms, min_time 1.24 ms
--------------------------------------

六性能統(tǒng)計(jì)

AX650N總算力10.8Tops@Int8，支持硬切分為三個(gè)獨(dú)立小核心或一個(gè)大核心的能力，即：

●NPU1，3.6Tops

●NPU3，10.8Tops

七

后續(xù)計(jì)劃

●嘗試部署基于Transformer網(wǎng)絡(luò)結(jié)構(gòu)的分割模型，敬請(qǐng)期待。

審核編輯：湯梓紅