[发明专利]一种基于概率图和ViT模型的图片分类方法

说明书

本发明涉及一种基于概率图和ViT模型的图片分类方法，属于计算机视觉图片分类技术领域，将多头注意力机制从概率论的角度进行建模，将多头注意力中的attention value看作隐变量，利用概率图模型的Explaining‑away Effects以及Transformer的层级结构，将attention logits层层传递，并将相邻层的值进行融合，促进不同头部之间的交互。本发明针对普通Vision Transformer模型中多头注意力机制头部参数的冗余问题，将多头注意力机制建模为概率图模型，将注意力值看作隐变量，促进不同注意力头部之间的交互。

技术领域

本发明涉及一种基于概率图和ViT模型的图片分类方法，属于计算机视觉图片分类技术领域。

背景技术

Transformer在自然语言处理领域取得了巨大的成功，激励了人们尝试将多头注意力机制引入主流框架为卷积神经网络的计算机视觉领域。相较于卷积神经网络，Transformer在捕捉图片全局信息方面有着巨大的优势，同时，Transformer的可并行化计算也促进了其在视觉领域的应用。目前Vision Transformer在计算机视觉的各类任务，如图片分类、目标检测和图片降噪等方面取得了令人瞩目的效果。但是大量研究发现，Transformer的核心多头注意力机制中的不同头之间存在参数冗余，严重影响了模型的整体性能。

发明内容

本发明的目的是提供一种基于概率图和ViT模型的图片分类方法，用来解决上述缺陷。

为了实现上述目的，本发明采用的技术方案是：

一种基于概率图和ViT模型的图片分类方法，包括以下步骤：

S1、将输入模型的图片进行分块，然后将每个图片块展平成一维向量，最后通过线性变换生成patch embedding；

S2、给每个patch embedding加上位置编码，补充位置信息；

S3、增加一个用于分类的Token，学习其他图片patch的整体信息；

S4、基于头部交互的Transformer Block，把attention values看作隐变量，利用概率图模型中的Explaining-away Effects以及Transformer的层级结构，将attentionlogits层层传递，并将相邻层的值进行融合，促进不同头部之间的交互；

S5、使用两层全连接层，将分类Token输入分类层，得到图片的分类结果。

本发明技术方案的进一步改进在于，所述S1的具体步骤为：

S11、将输入模型的图片进行分块、展平，具体操作为：

将图片patch的长宽均设置为P，即将图片数据H*W*C变换为

其中，N为一张图分割的patch数量，C为通道数，H为图片高度，W为图片宽度；

S12、将patch向量线性变换为patch embedding：

patch_embedding＝nn.Linear(patch_dim,dim)

其中，patch_dim为patch向量的维度，dim为patch embedding的维度。

本发明技术方案的进一步改进在于，所述S2的具体操作为：

pos_embedding＝nn.Parameter(torch.randn(1,num_patches+1,dim))

其中，pos_embedding为patch的位置编码，num_patches为patch的数量。