V-Tech | 为旌星图工具链部署SwinT模型,降低视觉Transformer计算成本

从Transformer到SwinT

Transformer算法最初是为自然语言处理任务设计的,因其注意力机制的独特设计,使得模型能够捕捉到全局和局部之间的复杂关系,因而计算机视觉领域一直尝试将其引入。

2020年,DETR算法成功将CNN与Transformer相结合,构建了一个端到端的目标检测模型。同年提出的ViT算法,创新性地将图像划分成一个个patch,并将每个patch展平为一个向量,使得图像数据转化为序列化数据,之后输入到Transformer模型中,实现了Transformer在图像分类任务中的应用。ViT算法相比于DETR算法,摒弃了以CNN网络为backbone的做法,从输入图像下手,将其转变为序列数据,基本无需改动Transformer便可实现图像分类网络。这种处理方式,也给后续许多算法提供了新思路。

但ViT依旧存在局限性,由于计算量与像素数量的平方成正比,对于高清图像,计算量陡然上升。

为解决上述问题,诞生了SwinT(Shifted Windows Transformer)。该算法的主要设计思路围绕以下两点展开:

模仿CNN算法的卷积操作,提出一种滑动窗口机制,通过不重叠的local-window和重叠的cross-window,将Transformer的全局注意力计算限制在一个大小固定的窗口之中。

加入下采样的层级设计,从而逐层扩大感受野范围,使得被限制在窗口范围内的注意力计算也能获取全局的特征信息。

由于引入了类似CNN卷积的操作,再加上多层级的下采样结构,SwinT既具备了较强的局部信息聚合能力,又保留了Transformer无需过多操作便能捕获远距离特征信息的特点,因此,算法表现极为出色,既能够作为Transformer类型算法在计算机视觉领域的通用backbone,又能在一定程度上能代替CNN。

SwinT算法原理

网络主体结构

SwinT网络结构,由Patch Partition和4个Stage组成。Patch Partition将输入图像转换成patch size为4的众多patch。每个Stage包含一个Linear Embdding或Patch Merging,以及若干STB模块(Swin Transformer Block)。在Stage1中先通过Linear Embedding将众多patch嵌入到高维向量中,即把通道数调整为C,之后送入STB。Stage2-3,均为先由Patch Merging进行下采样,缩小特征图的分辨率,扩大感受野范围,然后送入STB。

SwinT架构

Patch Merging

在CNN中,通常会使用卷积层或池化层进行下采样操作,而SwinT的下采样操作相对复杂,每次下采样均为两倍,因此在H/W方向上,均是间隔2选取元素,得到四个宽高为(H/2,W/2)的tensor,通过将这四个tensor在c通道方向堆叠,形成一个新的tensor后,经过Linear层,最终获得下采样输出。

Patch Merging流程

tensor维度变化

STB模块

STB模块是SwinT中最核心的部分,其在标准的Transformer自注意力基础上改进为窗口自注意力(W-MSA:window multi-head self-attention)和移动窗口自注意力(SW-MSA:shifted window multi-head self-attention)两种结构。这两种结构是成对使用的,其中:

W-MSA将每一个window范围当作独立的全局进行自注意力计算,从而降低计算量;

SW-MSA则把window偏移一半边长后计算window范围内的注意力,使得相邻window内的元素产生交互,进而实现特征信息融合。

这种通过window移位实现不相交window内数据聚合的方法,有效提升了各类计算机视觉任务中的计算效果。

STB结构

需要说明的是,SW-MSA部分由于对window进行了移位操作,使得计算量增加。以下图为例,左边进行的是W-MSA操作,此时只有4个window需要进行内部的注意力计算。在SW-MSA操作中(右图),window偏移了一定距离后,出现了9个大小不统一的window需要进行注意力计算,这势必会带来额外的计算开销。

对于移位引起的计算量增加,SwinT采用了cyclic shift和attention mask的方案来解决。

cyclic shift通过按照一定规律移动feature map,将window移位后产生的9个计算区域,重新拼接成4个区域,使得计算量和W-MSA保持一致。首先,将顶部数据(序号1-4)移动到底部,再把最左侧数据(序号5、9、13、1)移动到最右侧,得到result feature map。

result feature map中的4个区域,除左上角区域没有变化,其他区域都由原本相距很远的数据拼接而成,这部分数据本不应参与所处区域的注意力计算,故要对其进行屏蔽,即使用attention mask来掩盖住不属于本区域的数据,得到注意力结果之后再把移位区域平移回去。

cyclic shift+attention mask解决移位计算量增加问题

SwinT模型导出Onnx及量化

为旌天权TMNPU是一款专为深度学习算法优化加速而设计的硬件单元,其广泛兼容并支持当前主流公开的神经网络以及各类常见算子,涵盖了深度学习卷积神经网络和Transformer网络。为旌星图TM工具链针对为旌全自研天权NPU硬件平台开发,充分利用为旌天权TMNPU所特有的高集成度自定义指令集,从而确保深度学习网络模型能够得到高性能执行。

在模型部署过程中,为旌星图TM工具链要求待部署的模型格式为Onnx,即非Onnx格式的原模型需要先进行转换,下文以pytorch模型为例。

Onnx模型导出

pytorch模型一般以源码+参数文件的形式保存。作为示例,本文一同使用的torchvision中的模型也是以上述形式保存。

环境准备

首先是确保pytorch和torchvision已安装且torchvision的版本在0.13以上。请注意python、torchvison和pytorch之间的版本存在约束关系,具体请参考https://github.com/pytorch/vision。本文在python3.8和cuda11.7的环境中安装pytorch2.0.0和torchvison0.15.1,使用如下命令:

pip install torch==2.0.0 torchvision==0.15.1

如软硬件环境不同,请参考https://pytorch.org/get-started/previous-versions/选择命令安装。

torch模型准备

使用torchvision中的SwinT模型,注意需要加载预训练参数。命令如下:

import torch

import torchvision

weigths = torchvision.models.Swin_T_Weights

model = torchvision.models.swin_t(weights=weigths)

model.eval()

转换Onnx

从torchvision官网的介绍中我们可以看到该模型是一个基于imagenet数据集的分类模型,该模型只有一个输入和一个输出。

设置输入名称为"x",输出名称为“y”,输入的shape为[1,3,224,224];设置模型的存储文件名为swin_t.onnx,opset_verion为13。综上,使用如下命令即可导出Onnx模型:

shape = [1,3,224,224]

x = torch.randn(size=shape)

save_path = 'swin_t.onnx'

input_names = ['x']

output_names = ['y']

opset_version=13

torch.onnx.export(

model, x, save_path, input_names=input_names, output_names=output_names, opset_version=opset_version

)

上述代码运行之后,得到swin_t.onnx。该文件可用netron工具打开。

更多具体转换教程可以查看此网址:https://pytorch.org/docs/2.0/onnx.html。

模型量化

模型清洗

我们把开发环境切换为旌星图TM工具链的开发环境。

此前转换出的Onnx模型中并不能直接使用,因为工具链不支持动态算子(比如shape)、分支和循环,另外本身模型中还有很多待优化的地方。使用onnxsim可以进行常量折叠、算子融合、死代码消除等操作。onnxsim的命令如下:

onnxsim swin_t.onnx swin_t_new.onnx

得到swin_t_new.onnx之后,推荐再使用OnnxConvertTool命令进行模型清洗,该命令可以将Onnx模型中一些不支持的结构(比如不对称pad的卷积和池化算子)转为支持的结构,也会将模型的一些算子和参数替换为等效的算子和参数。该命令如下:

OnnxConvertTool --input_model swin_t_new.onnx --output_model swin_t_washed.onnx

使用该命令,我们最终得到swin_t_washed.onnx,作为量化用的模型。

PTQ

在运行ptq之前,我们需要准备数据集。下文仅作演示,我们从imagenet数据集中随机抽取了20张图片作为校准数据集,并放在名为imagenet_demo的文件夹下。

在imagenet_demo的同级目录下,新建config_ptq.yml,内容如下:

model:

onnx_model: ./swin_t_washed.onnx

work_mode: PTQ

quant:

activation:

observer: MinMaxObserver

dtype: int8

symmetry: True

weight:

observer: MinMaxObserver

dtype: int8

symmetry: True

per_channel: True

dataset:

calibrate:

root: ./imagenet_demo/

calibrate_num_sampler: 20

batch_size: 1

transform:

resize: [256, 256]

crop: [224, 224]

out_dir: ./swint_quant_out

运行如下命令启动ptq量化:

StatlasQuant --quant_cfg config_ptq.yml

量化完成之后,在当前目录下找到swint_quant_out文件夹,可以看到该文件夹产生了一个deploy的Onnx模型以及量化参数表,这两个文件可供后续的编译器使用。

$ tree swint_quant_out/

swint_quant_out/

├── swin_t_washed_deploy_model.onnx

└── swin_t_washed_quant_param.yaml

经过以上步骤,原模型转Onnx并量化的工作完成。转换得到的量化参数表和Onnx模型会继续进入编译器,通过编译生成megrez文件,该文件可以通过runtime工具在芯片上运行该模型。

关于编译器与runtime工具我们将在后续的文章中详细介绍。

小结

SwinT在机器视觉领域中表现出了出色的性能,通过将输入图片分割成小块,并进行自注意力计算,实现了高效的图像特征提取。

除本文重点介绍的SwinT外,为旌全自研天权NPU广泛支持已公开的主流神经网络模型以及各类常见算子。配套的为旌星图TM工具链旨在帮助用户将神经网络模型高效、快速地部署到为旌天权TMNPU上进行加速处理,有效缩短客户的开发周期。通过整合端到端的解决方案,企业可以更加专注于业务逻辑与模型创新,而无需过多关注底层基础设施和运维难题,从而加速AI技术驱动的业务转型和升级。


夕夕海 » V-Tech | 为旌星图工具链部署SwinT模型,降低视觉Transformer计算成本

发表回复