[发明专利]一种基于注意力图的视觉Transformer模型剪枝方法

说明书

本发明公开了一种基于注意力图的视觉Transformer模型剪枝方法，应用于机器视觉推理系统，包括如下步骤：在机器视觉推理系统中，通过数据训练网络对ViT模型执行若干轮初始训练，以生成完整的注意力图；计算注意力图的信息熵，根据计算的信息熵大小对注意力头进行剪枝操作；移除被剪枝注意力头关联的各权重参数，以获得新ViT模型；重新微调新ViT模型的参数；通过对多头注意力模块进行剪枝，删除具有高不确定性的特征图和相应的注意力头，以减少ViT模型的参数和复杂性，降低ViT模型的计算复杂度和参数量，能够缩减ViT模型大小，最终达到在ViT模型性能有限损失的情况下实现ViT模型的轻量化。

技术领域

本发明属于神经网络轻量化技术领域，具体涉及一种基于注意力图的视觉Transformer模型剪枝方法。

背景技术

Transformer是一种主要基于自注意机制的深度神经网络，应用于自然语言处理领域，视觉Transformer模型简称ViT模型，Transformer具有强大的长程依赖关系的建模能力，已经在各种视觉任务中取得了令人瞩目的成功，然而，Transformer模型的巨大计算量和内存消耗是其固有问题，使其无法在资源有限的边缘端计算设备上成功部署并投入使用，剪枝是一种有效降低神经网络推理成本的常用方法，已广泛应用于计算机视觉和自然语言处理应用中。

基于注意力图的模型剪枝方法可用于将神经网络模型部署在低功耗、计算资源受限的嵌入式机器视觉推理系统中，包括基于图形处理器加速的嵌入式计算板和基于神经网络处理器，该类系统一般只能提供相当于高性能GPU不到20%的计算资源。

剪枝操作一般可以分为非结构化剪枝和结构化剪枝两大类，具体来说，非结构化剪枝在特定标准下删除单个不重要的权重，非结构化剪枝属于精细范式，对精度有少许损害，实际加速需要特殊的硬件设计，结构化剪枝移除了模型的整个子结构，例如通道和注意力头，已经有一些工作通过减少计算图像编码块的数量修剪ViT，Tang等人开发了一种自上而下的图像块剪枝方法，该方法基于预训练模型的重建误差去除冗余图像块，Xu等人基于结构维持的图像编码块选择和慢－快结合的更新策略以完全利用整个空间结构；上述方法虽然可以节省计算成本，但不能降低推理复杂性和减小模型大小，为此我们提出一种基于注意力图的视觉Transformer模型剪枝方法。

发明内容

本发明的目的在于提供一种基于注意力图的视觉Transformer模型剪枝方法，以解决上述背景技术中提出的问题。

为实现上述目的，本发明提供如下技术方案：一种基于注意力图的视觉Transformer模型剪枝方法，应用于机器视觉推理系统，包括如下步骤：

步骤A、在机器视觉推理系统中，通过数据训练网络对ViT模型执行若干轮初始训练，以生成完整的注意力图；

步骤B、计算注意力图的信息熵，根据计算的信息熵大小对注意力头进行剪枝操作,度量注意力图的不确定性；

步骤C、移除被剪枝注意力头关联的各权重参数，以获得新ViT模型；

步骤D、重新微调新ViT模型的参数。

优选的，在所述的步骤A中，ViT模型将输入图像拆分为N个图像块，并对每个图像块附加一个类编码，然后将附加类编码的N个图像块馈送到与普通Transformer类似的编码器中，形成N个图像编码块。

优选的，所述的步骤A，包括如下步骤：

A1、在ViT模型训练的初始阶段，ViT模型没有学习到有用的信息，此时注意力图是无序的，且注意力图具有大的信息熵；

A2、在ViT模型经过若干轮的初始训练，ViT模型学习到基本信息，并开始呈现出一定的模式；

A3、在ViT模型训练的最终阶段，当ViT模型收敛时，每个注意力头都获得了注意力图，此时重要的图像编码块会受到注意力头的高度关注，信息熵会降低，所有注意力图都是一个训练轮次的平均结果。