01

背景

今年來以ChatGPT為代表的大模型的驚艷效果，讓AI行業迎來了新的動力。各種AIGC的應用接踵而至。我們知道類似ChatGPT的大模型，其核心網絡結構均基于Google 2017年的論文提出的Transformer的論文《Attention Is All You Need》。在計算機視覺建模一直由卷積神經網絡（CNN）主導，基于Transformer結構的網絡模型長時間停留在各大頂會“刷榜”階段，真正大規模落地并不突出。直到ICCV 2021的最佳論文《Swin Transformer》才達到了準確率和性能雙佳的效果。

但是到目前為止，類似Swin Transformer的視覺類Transformer網絡模型大多數還是部署在云端服務器上，原因是GPU對于MHA結構計算支持更友好，反而邊緣側/端側AI芯片由于其DSA架構限制，為了保證CNN結構的模型效率更好，基本上對MHA結構沒有過多性能優化，甚至需要修改網絡結構才能勉強部署。這也間接限制了算法工程師在邊緣計算應用上進一步發揮Transformer網絡的想象力。

今年3月，愛芯元智發布了新一代產品AX650N，內置了其自主研發的第三代神經網絡單元，進一步提升了最新AI算法模型的部署能力，可幫助用戶在智慧城市，智慧教育，智能制造等領域發揮更大的價值。最近我通過正式渠道有幸拿到了一塊AX650N Demo板進行嘗鮮體驗。

本文的目的是簡單介紹基于AX650N Demo配套的新一代AI工具鏈如何優雅地將Swin Transformer模型部署到AX650N Demo板上，希望能給算法工程師們在Transformer網路部署落地上提供一種新的思路和途徑。

02

Swin Transformer

The architecture of a Swin Transformer

目前Transformer應用到圖像領域主要有兩大挑戰：

視覺實體變化大，在不同場景下視覺Transformer性能未必很好；

圖像分辨率高，像素點多，Transformer基于全局自注意力的計算導致計算量較大。

2.1 原理

針對上述兩個問題，微軟在《Swin Transformer》的論文中提出了一種包含滑窗操作。其中滑窗操作包括不重疊的local window，和重疊的cross-window。將注意力計算限制在一個窗口中，一方面能引入CNN卷積操作的局部性，另一方面能節省計算量。在各大圖像任務上，Swin Transformer都具有很好的性能。

2.2 分析

相比常見CNN網絡模型，其實也就是新增了MHA（Multi Head Attention）的關鍵算子

LayerNormalization

Matmul

GELU

量化

LN、GELU、Matmul存在掉點風險

計算效率

占比最大的計算操作由Conv變成Matmul，因此要求硬件平臺MatMul計算能力強

03

模型轉換

Pulsar2介紹

Pulsar2（暫定名）是我們的新一代AI工具鏈，在吸取上一代工具鏈Pulsar的優秀行業經驗和不足之處的反思后進行的重構，依然包含“模型轉換、離線量化、模型編譯、異構調度”四合一功能，進一步強化的網絡模型快速、高效的部署需求。在針對第三NPU架構進行了深度定制優化的同時，也擴展了算子&模型支持的能力及范圍，對Transformer結構的網絡也有較好的支持。

pulsar2 deploy pipeline

3.1 模型下載

從Swin Transformer的官方倉庫獲取模型，由于是基于PyTorch訓練，導出的是原始的pth模型格式，而對于部署的同學而言，更喜歡使用ONNX模型進行后續的產品落地，為了方便測試，我們提供該模型的ONNX版本導出腳本，降低模型獲取門檻，便于之前不熟悉的同學直接掌握其中的關鍵操作。

import onnx
import torch
import requests
from onnxsim import simplify
from PIL import Image
from transformers import AutoFeatureExtractor, SwinForImageClassification


def download_swin_model(model_name):
  prefix = "microsoft"
  model_id = f"{prefix}/{model_name}" # google/vit-base-patch16-384


  url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
  image = Image.open(requests.get(url, stream=True).raw)
  feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
  model = SwinForImageClassification.from_pretrained(model_id)
  inputs = feature_extractor(images=image, return_tensors="pt")
  outputs = model(**inputs)
  logits = outputs.logits
  # model predicts one of the 1000 ImageNet classes
  predicted_class_idx = logits.argmax(-1).item()
  print("Predicted class:", model.config.id2label[predicted_class_idx])


  # export
  model_path = f"{model_name}.onnx"
  torch.onnx.export(
    model,
    tuple(inputs.values()),
    f=model_path,
    do_constant_folding=True,
    opset_version=13,
    input_names=["input"],
    output_names=["output"]
  )


  # simplify
  model = onnx.load(model_path)
  model_simp, check = simplify(model)
  assert check, "Simplified ONNX model could not be validated"
  simp_path = f"{model_name}_sim.onnx"
  onnx.save(model_simp, simp_path)


def main():
  download_swin_model(model_name="swin-tiny-patch4-window7-224") # microsoft/swin-tiny-patch4-window7-224


if __name__ == "__main__":
  main()

3.2 模型編譯

Pulsar2為了提升用戶使用體驗，降低Pulsar客戶遷移的學習成本，基本上延續了原有風格，包括Docker環境安裝、命令行指令、配置文件修改參數、仿真功能等。同時針對編譯速度慢的痛點，進行了大幅度優化，模型編譯的耗時相比第一代工具鏈平均降低了一個數量級（分鐘->秒）。

$ pulsar2 build --input model/swin-t.onnx --output_dir output --config config/swin-t.json --target_hardware=AX650
32 File(s) Loaded.
[10:22:36] AX Quantization Config Refine Pass Running ... Finished.
[10:22:36] AX Quantization Fusion Pass Running ...    Finished.
[10:22:36] AX Quantize Simplify Pass Running ...     Finished.
[10:22:36] AX Parameter Quantization Pass Running ...   Finished.
Calibration Progress(Phase 1): 100%|████████| 32/32 [00:08<00:00, ?3.92it/s]
Finished.
[10:22:45] AX Passive Parameter Quantization Running ... ?Finished.
[10:22:45] AX Parameter Baking Pass Running ... ? ? ? ? ? Finished.
[10:22:45] AX Refine Int Parameter pass Running ... ? ? ? Finished.
Network Quantization Finished.
quant.axmodel export success: output/quant/quant_axmodel.onnx
Building native ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
......
2023-04-13 10:23:07.109 | INFO ? ? | yasched.test_onepass1475 - max_cycle = 6689562
2023-04-13 10:23:25.765 | INFO ? ? | yamain.command.build832 - fuse 1 subgraph(s)

從編譯log中我們大致看出，計算圖優化、PTQ量化、離線編譯總共耗時只需50秒。然后我們來看一下大家比較關心的MHA結構變成了什么樣子：

MHA ONNX原始結構

MHA由工具鏈進行圖優化之后的quant.axmodel結構

3.3 仿真運行

在這一代工具鏈，我們提供更方便的pulsar2-run-helper的插件，可以模擬NPU計算流程，方便提前獲得上板運行結果。（請大家記住仿真運行的結果，后續章節將與上板實際部署的推理結果進行比對）

python3 cli_classification.py --post_processing --axmodel_path models/swin-t.axmodel --intermediate_path sim_outputs/0
[I] The following are the predicted score index pair.
[I] 2.6688, 285
[I] 1.9528, 223
[I] 1.8877, 279
[I] 1.8877, 332
[I] 1.8226, 282

04

上板部署

AX650N Demo板的BSP上已經預裝了NPU模型測試需要的工具

/root # sample_npu_classification -m swin-t.axmodel -i cat.jpg -r 100
--------------------------------------
model file : swin-t.axmodel
image file : cat.jpg
img_h, img_w : 224 224
--------------------------------------
Engine creating handle is done.
Engine creating context is done.
Engine get io info is done.
Engine alloc io is done.
Engine push input is done.
--------------------------------------
2.6688, 285
1.9528, 223
1.8877, 332
1.8877, 279
1.8226, 282
--------------------------------------
Repeat 100 times, avg time 8.64 ms, max_time 8.65 ms, min_time 8.64 ms
--------------------------------------

對比上一章節的仿真結果，完全一致。

4.1 算力分配

AX650N的10.8Tops@Int8的算力其實是可分配的，上述內容中，按照默認的編譯選項，其實只發揮了一部分算力（3.6Tops@Int8）。我們來看看滿算力下的耗時表現如何呢？

/root # ax_run_model -m swin-t-npu3.axmodel -r 100
Run AxModel:
   model: swin-t-npu3.axmodel
    type: NPU3
    vnpu: Disable
  affinity: 0b001
   repeat: 100
   warmup: 1
   batch: 1
  tool ver: 1.0.0
 ------------------------------------------------------
 min =  3.769 ms  max =  3.805 ms  avg =  3.778 ms
 ------------------------------------------------------
/root #
/root # sample_npu_classification -m swin-t-npu3.axmodel -i cat.jpg -r 100
--------------------------------------
model file : swin-t-npu3.axmodel
image file : cat.jpg
img_h, img_w : 224 224
--------------------------------------
Engine creating handle is done.
Engine creating context is done.
Engine get io info is done.
Engine alloc io is done.
Engine push input is done.
--------------------------------------
2.6688, 285
1.9528, 223
1.8877, 332
1.8877, 279
1.8226, 282
--------------------------------------
Repeat 100 times, avg time 3.78 ms, max_time 3.79 ms, min_time 3.77 ms
--------------------------------------

05

性能統計

NPU工具鏈的性能優化是個長期堅持的過程，最新版本的性能數據會更優秀。

審核編輯：劉清