如何优雅地将Swin Transformer模型部署到AX650N Demo板上？

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工具链的性能优化是个长期坚持的过程，最新版本的性能数据会更优秀。

审核编辑：刘清