基于AX650N部署EfficientViT

一背景

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

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

二EfficientViT

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

● 論文：

https://arxiv.org/abs/2305.070271

● Github鏈接：

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

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

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

另外考慮到BN可以與FC和Conv在推理時融合以實現(xiàn)加速。因此使用Batch Normalization替換Layer Normalization，同時在大尺度下層數(shù)更少，并在每個stage用了小于2的寬度擴展系數(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解碼。強大的性能可以讓AX650N幫助用戶在智慧城市，智慧教育，智能制造等領(lǐng)域發(fā)揮更大的價值。

四模型轉(zhuǎn)換

本文以EfficientViT-M5為例。

4.1 Pulsar2

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

npu-software-arch

4.2 模型下載

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

下載安裝

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

●獲取PyTorch模型

雖然官方倉庫中已經(jīng)提供了ONNX模型，但是我們發(fā)現(xiàn)其ONNX模型的Batch Size=16，并不適合普通的端側(cè)芯片進行評估，因此這里選擇從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)化計算圖

導(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)化、離線量化、編譯、對分功能。整個過程耗時只需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

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

圖優(yōu)化-前

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

五上板部署

5.1 AX-Samples

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

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

5.2 運行

運行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)計

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

●NPU1，3.6Tops

●NPU3，10.8Tops

七

后續(xù)計劃

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

審核編輯：湯梓紅