如何优雅地将Swin Transformer模型部署到AX650N Demo板上的完整教程

01 背景

近年来，以ChatGPT为代表的大模型展现出的惊艳效果，为AI行业注入了全新动力，各类AIGC应用层出不穷。这些大模型的核心网络结构均基于Google 2017年提出的Transformer论文《Attention Is All You Need》。在计算机视觉领域，长期由卷积神经网络（CNN）主导，但基于Transformer结构的网络模型长时间停留在各大顶会“刷榜”阶段，真正大规模的落地并不突出。直到ICCV 2021的最佳论文《Swin Transformer》才实现了准确率与性能的双重突破。

然而，截至目前，类似Swin Transformer的视觉类Transformer网络模型大多仍部署在云端服务器上。原因在于GPU对MHA（多头注意力）结构的计算支持更为友好，而边缘侧/端侧AI芯片因其DSA架构限制，为确保CNN结构模型的效率，对MHA结构缺乏充分的性能优化，甚至需要修改网络结构才能勉强部署。这间接限制了算法工程师在边缘计算应用中进一步发挥Transformer网络的想象力。

2023年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网络模型，Swin Transformer主要新增了MHA（Multi Head Attention）关键算子，具体包括：

• LayerNormalization
• Matmul
• GELU

量化方面：LN、GELU、Matmul存在掉点风险。

计算效率方面：占比最大的计算操作由Conv变成了Matmul，因此要求硬件平台具备较强的MatMul计算能力。

小提示：在实际部署时，需特别关注LN、GELU、Matmul的量化精度，建议进行充分的校准和验证。

常见问题：

Q：为什么Swin Transformer要使用滑窗操作？

A：滑窗操作通过局部窗口限制注意力计算范围，既降低了计算复杂度，又保留了CNN的局部归纳偏置，使得模型在图像任务中表现更优。

Q：Swin Transformer相比ViT有哪些优势？

A：Swin Transformer通过层级式结构和滑窗注意力，能够处理高分辨率图像，并具有线性计算复杂度，更适合实际场景。

03 模型转换

3.1 Pulsar2介绍

Pulsar2（暂定名）是爱芯元智的新一代AI工具链。在吸取上一代工具链Pulsar的优秀经验和不足的基础上进行了重构，依然包含“模型转换、离线量化、模型编译、异构调度”四合一功能，进一步强化了网络模型的快速、高效部署需求。针对第三代NPU架构进行了深度定制优化的同时，扩展了算子与模型支持的能力及范围，对Transformer结构的网络也有较好的支持。

pulsar2 deploy pipeline

小提示：使用Pulsar2前请确保已安装Docker环境，并拉取最新版Pulsar2镜像。

3.2 模型下载

从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.sa ve(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()

小提示：运行脚本前请安装依赖库：pip install transformers onnx onnx-simplifier pillow requests

3.3 模型编译

Pulsar2为提升用户使用体验，降低学习成本，基本延续了原有风格，包括Docker环境安装、命令行指令、配置文件参数修改、仿真功能等。同时针对编译速度慢的痛点进行了大幅优化，模型编译的耗时相比第一代工具链平均降低了一个数量级（从分钟级降至秒级）。

$ pulsar2 build --input model/swin-t.onnx --output_dir output --config config/swin-t.json --target_hardware=AX6503

2 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)

从编译日志可以看出，计算图优化、PTQ量化、离线编译总共仅耗时50秒。下面我们来看一下大家比较关心的MHA结构在编译后变成了什么样子：

MHA ONNX原始结构

MHA由工具链进行图优化之后的quant.axmodel结构

常见问题：

Q：编译过程中间出现量化校准失败怎么办？

A：请检查校准数据集是否正确配置，建议使用与模型训练时相似的图片数据，确保覆盖常见场景。

Q：生成的axmodel文件如何查看结构？

A：可以使用Netron工具打开quant.axmodel文件（该文件为ONNX格式），直观查看网络结构和参数。

3.4 仿真运行

新一代工具链提供了更便捷的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

小提示：仿真运行是通过NPU指令模拟器在CPU上模拟执行，可用于快速验证模型精度和性能预估，但实际上板性能可能略有差异。

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, a vg 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
  a vg =   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, a vg time 3.78 ms, max_time 3.79 ms, min_time 3.77 ms
--------------------------------------

常见问题：

Q：如何切换算力模式？

A：在模型编译时，通过配置文件中设置npu_mode或使用不同编译选项生成对应算力版本的axmodel文件即可。具体可参考Pulsar2用户手册。

Q：上板推理结果与仿真不一致怎么办？

A：请检查模型编译时是否开启了相同的量化参数，以及校准数据集是否一致。另外，检查输入图像的预处理方式（如归一化、尺寸调整）是否与训练时一致。

05 性能统计

NPU工具链的性能优化是一个长期坚持的过程，最新版本的性能数据会进一步提升。本次部署验证了Swin Transformer在AX650N上的高效运行，为边缘端Transformer模型的落地提供了可靠参考。

小提示：若需要更高帧率，可在编译时开启模型优化选项（如算子融合、内存复用等），同时合理选择算力模式以平衡功耗与性能。

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：如何优雅地将Swin Transformer模型部署到AX650N Demo板上的完整教程要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题