模型推理:Swin-Transformer图像分类

模型推理:Swin-Transformer图像分类

本实验介绍了基于寒武纪 MagicMind 的模型推理—Swin-Transformer 图像分类的实践过程
难易程度: 中级|实验人次:2371

一、内容和目标

 

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
  1. 下载 Swin-Transformer 预训练模型,并将其放置于pytorch_swin_transformer_inference目录下,本实验的 patch.sh已自动操作。
  1. 搭建环境,本实验的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:修改AdaptiveAvgPool1dnn.AdaptiveAvgPool2d
  1. 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

 

申 请 试 用