如何优雅地将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

性能统计

算力	耗时（ms）	帧率（fps）
3.6Tops@Int8	8.64	115
10.8Tops@Int8	3.77	265

NPU工具链的性能优化是个长期坚持的过程，最新版本的性能数据会更优秀。

审核编辑：刘清