一、内容和目标

 

1. 实验内容

1. 本实验主要介绍基于寒武纪 MLU370 MagicMind 平台的Swin-Transformer (Pytorch, Python, FP32) 图像分类推理应用的开发方法。

2. 编写自定义算子 Plugin Roll 和 Plugin ReLU，生成含有自定义算子的 PyTorch 模型。

3. 基于 Swin-Transformer 分类网络和寒武纪 MLU370 MagicMind 平台，您可以读取本地图像数据作为输入，对图像进行分类。

 

2. 实验目标

1. 掌握编写 PluginOp 的方法，生成含有自定义算子的 PyTorch 模型。

2. 掌握使用寒武纪 MLU370 MagicMind 平台进行 AI 模型推理的基本方法。

3. 理解 Swin-Transformer 模型的整体网络结构及其开发调试细节。

 

二、平台介绍

• 硬件：寒武纪 MLU370 AI 加速卡

• 框架：Pytorch 1.6 、MagicMind 0.14.0

• 系统环境：寒武纪云平台

 

三、网络结构

Swin-Transformer 在多项视觉任务中（分类、检测、分割）都有着较好的表现，对于分类任务，Swin-Transformer 在网络最后面会接上一个 Layer Norm 层、全局池化层以及全连接层得到分类输出。Swin-Transformer 针对不同大小（Swin-T、Swin-S、Swin-B、Swin-L）的网络整体架构不变，但会根据 yaml 文件中定义的 EMBED_DIM 和 DEPTHS 参数，对每个子模块进行不同通道数和深度的缩放。

以swinv2_tiny为例，其网络结构主要由以下几部分组成:

EMBED_DIM: 96 表示第一个stage中隐藏层的通道数。

DEPTHS: [ 2, 2, 6, 2 ] 表示每个 stage 中 Swin Transformer Block 的个数。

• swinv2_tiny的网络结构如下图所示:

 

• Swin Transformer Block 结构如下图所示：

 

其中 W-MSA 和 SW-MSA 分别是包含常规窗和滑动窗的多头自注意力结构。上述结构图来源于 Swin-Transformer论文。

 

 

四、实验流程

本实验按照下述流程进行展开：

 

五、基于CNNL实现编写自定义算子Plugin Roll

MagicMind除了提供大量的内置算子（如conv、pool等），还提供自定义算子开发功能，来满足用户的业务需求和性能需求。当MagicMind Parser 解析原生 PyTorch 模型某个算子失败时，可编写自定义PluginOp，替换 PyTorch 原生网络中对应算子，以使模型能被 MagicMind Parser 正常解析。

编写自定义算子的一个核心任务是实现算子本身的功能，在寒武纪软件栈中，实现算子的功能有两种方式，分别是BANG C实现和CNNL实现。其中BANG C是CNNL底层的实现，CNNL又是对 BANG C 做了进一步的封装和优化。

基于CNNL实现编写自定义算子 Plugin Roll主要包含四个步骤：注册自定义算子、实现算子Kernel、注册算子Kernel、通过parser使用自定义算子。其中第4步又包含3个子步骤，即编译生成 MagicMind 自定义算子插件、为 MagicMind Pytorch Parser 添加MagicMind 自定义算子插件、PluginOp替换原生网络中对应算子。

 

六、基于BANG C实现编写自定义算子Plugin ReLU

BANG C是一门为MLU硬件编程的语言，BANG C语言基于C/C++，并在其基础上增加了BANG异构并行编程模型必须的语法特性、计算原语、数据类型和内建变量支持。同时也禁用了一些不适合异构编程环境的C/C++特性。

BANG 异构并行编程模型利用 CPU 和 MLU 协同计算，实现了 CPU 和 MLU 的优势互补。在BANG 异构并行编程模型中，CPU作为主机侧的控制设备，用于完成复杂的控制和任务调度；而设备侧的 MLU 则用于大规模并行计算和领域相关的计算任务。设备侧的入口函数就是Kernel函数，Kernel函数中可以使用BANG C的语法特性实现算子功能。

基于BANG C实现编写自定义算子 Plugin ReLU主要包含四个步骤：注册自定义算子、实现算子Kernel、注册算子Kernel、通过parser使用自定义算子。其中第4步又包含3个子步骤，即编译生成 MagicMind 自定义算子插件、为 MagicMind Pytorch Parser 添加MagicMind 自定义算子插件、PluginOp替换原生网络中对应算子。

相比于Plugin Roll直接调用CNNL接口，Plugin ReLU调用BANG C接口时，需要给出BANG C源码实现，且源码参与后续步骤的编译过程。

 

七、编写推理应用

编写推理应用包含五部分内容，分别是工程准备、编译PluginOp、修改Swin-Transformer源码、模型生成，模型推理。

 

STEP 1. 工程准备

1. 下载原始 Swin-Transformer 工程，本实验已经提供工程环境，Swin-Transformer源代码在src/目录下。对应官方commit id为：d19503d7fbed704792a5e5a3a5ee36f9357d26c1

2. 下载 Swin-Transformer 预训练模型，并将其放置于pytorch_swin_transformer_inference目录下，本实验的 patch.sh已自动操作。

3. 搭建环境，本实验的sh已自动操作。

 

STEP 2. 编译PluginOp

1. 编译Plugin ReLU。

2. 编译Plugin Roll。

 

STEP 3. 修改Swin-Transformer源码

1. src/config.py：注释掉LOCAL_RANK = args.local_rank代码，该代码为分布式训练相关。
2. src/models/swin_transformer_v2.py：修改AdaptiveAvgPool1d为nn.AdaptiveAvgPool2d。

3. src/models/swin_transformer_v2.py：使用ops.MyPlugin.PluginReLU替换nn.ReLU；使用torch.ops.MyPlugin.cnnlRoll替换torch.roll。

可使用为源码打patch方法完成上述操作。

 

STEP 4. 模型生成

1. 构建融合图。

2. 加载.so文件实现libtorch的算子注册。

3. Network和BuildConfig配置。

4. 生成MagicMind模型。

 

STEP 5. 模型推理

1. 图像前处理：使用 pillow 读取测试图片，预处理步骤包含Resize、center_crop与归一化。

2. 利用 MagicMind 进行推理，推理流程如下。

    3. 推理结果说明：

（1）在推理时，官方的Swin-Transformer网络末尾无softmax层。实验中为了体现每一类的分类概率，在MLU返回推理结果后加上该层。

（2）MagicMind 推理完成后，需要在 CPU 上进行softmax（可选），top1和top5的计算。

（3）推理完成后，在屏幕上会输出对每张图片预测的top1和top5结果。top1_cls、top5_cls表示图片所属的前1和前5类别，top1_prob、top5_prob表示图片所属前1和前5类别对应的概率。

（4）本次任务为1000分类，其中数字0-6 表示"鱼"，数字151-268 表示"狗"。更多详细类别可参考label.txt文件。

（5）为了显示表示模型运行到PluginOp，推理成功后，还会输出"plugin relu success" 和 "plugin roll success"。

 

八、相关链接

实验代码仓库：https://gitee.com/cambricon/practices

Modelzoo仓库：https://gitee.com/cambricon/modelzoo