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工具链的性能优化是个长期坚持的过程,最新版本的性能数据会更优秀。
审核编辑:刘清
-
gpu
+关注
关注
28文章
4673浏览量
128558 -
cnn
+关注
关注
3文章
350浏览量
22128 -
卷积神经网络
+关注
关注
4文章
359浏览量
11829 -
ChatGPT
+关注
关注
29文章
1546浏览量
7346
原文标题:爱芯分享 | 基于AX650N部署Swin Transformer
文章出处:【微信号:爱芯元智AXERA,微信公众号:爱芯元智AXERA】欢迎添加关注!文章转载请注明出处。
发布评论请先 登录
相关推荐
评论