引用本文

宋燕, 王勇. 多阶段注意力胶囊网络的图像分类. 自动化学报, 2024, 50(9): 1804−1817 doi: 10.16383/j.aas.c210012

Song Yan, Wang Yong. Multi-stage attention-based capsule networks for image classification. Acta Automatica Sinica, 2024, 50(9): 1804−1817 doi: 10.16383/j.aas.c210012

http://www.aas.net.cn/cn/article/doi/10.16383/j.aas.c210012

关键词

图像分类，胶囊网络，注意力机制，多阶段，鲁棒性 

摘要

针对传统的胶囊网络(Capsule network, CapsNet)特征提取不充分的问题, 提出一种图像分类的多阶段注意力胶囊网络模型. 首先, 在卷积层对低层特征和高层特征分别采用注意力(Spatial attention, SA)和通道注意力(Channel attention, CA)来提取有效特征; 然后, 提出基于向量的注意力(Vector attention, VA)机制作用于动态路由层, 增加对重要胶囊的关注, 进而提高低层胶囊对高层胶囊预测的准确性; 最后, 在五个公共数据集上进行图像分类的对比实验. 结果表明, 所提出的CapsNet模型在分类精度和鲁棒性上优于其他胶囊网络模型, 在仿射变换图像重构方面也表现良好.

文章导读

图像分类是指根据图片中的信息将图片划分到某一类别, 因此对图像进行特征信息提取是图像分类的重要研究内容. 传统的图像分类主要采用机器学习方法来提取特征, 随着深度学习的不断发展, 各种深度学习算法逐渐应用到图像分类当中. 2012年, AlexNet[1]神经网络在图像分类效果上超越了传统方法, 在AlexNet之后, 涌现出一系列改进的卷积神经网络(Convolutional neural network, CNN)模型[2-4], 不断地提高分类精度. 

然而, CNN的模型也存在一些缺陷. 首先, CNN的池化层会导致大量有价值的特征信息丢失, 从而对分类精度产生影响. 其次, 由于CNN对位置信息不敏感, 这将导致CNN对物体之间的空间关系的识别能力不强[5]. 随后提出的胶囊网络[6]则能够较好地处理上述问题, 具体地, 胶囊网络摒弃了CNN的池化层, 保留了大量的图片信息, 这使得胶囊网络运用较少的训练数据就能达到理想的效果. 此外, 胶囊网络是部分对整体的预测, 在预测的过程中能够较好地保留特征的姿态, 如位置、大小、方向等信息, 这使得胶囊网络不仅能够进行更加精确的分类, 还能够有效地识别出经过仿射变换等一系列空间变换的图像. 

近年来, 胶囊网络成为图像领域的一大研究热点. Sabour等[6]首先提出胶囊网络并且应用到图像分类任务, 作者基于公共数据集研究了胶囊网络的图像识别能力, 实验结果表明胶囊网络在图像分类的问题上可以成功地超越CNN. 然而, 由于胶囊网络在计算和内存负载方面的代价较高, 所以该胶囊网络的结构相对较浅, 主要适用于简单数据集, 在处理复杂数据时表现不佳. 随后, Deliege等[7]提出一种名为HitNet的深度学习网络, 该网络的核心思想是使用由胶囊组成的“Hit-or-Miss”层, 假定给定类的所有图像都具有各类特有的特征, 当给定一个输入图像, 强制HitNet瞄准真实类的胶囊所在的特定空间的中心, 而其他类的胶囊则被发送到远离各自特征空间中心的地方. 虽然该方法的收敛速度有大幅度提升, 但是复杂数据集上的精度有所降低. 针对该问题, 文献[8]提出多种改进的胶囊网络, 例如堆叠更多胶囊层、增加初始胶囊的数量、增加卷积层的层数或者更换其他激活函数等. 然而, 在复杂数据集上, 改进的胶囊网络在分类精度上没有明显的提升. 文献[9]中将原始胶囊网络中用向量表示的胶囊替换为矩阵表示, 同时将动态路由中的聚类算法思想换成高斯混合模型(Gaussian mixture model, GMM). 实验结果表明该模型仅在smallNORB数据集上有较小的提升, 复杂数据集上的效果依然不尽如人意. 

在注意力胶囊网络的研究方面, 相对于注意力卷积神经网络丰富的研究成果而言, 还有待进一步深入开展. 文献[10]在胶囊网络的卷积层中, 针对低层特征添加空间注意力机制, 虽然有效提取了特征之间的空间位置信息, 但缺乏对高层特征所描述的重要语义信息的特别关注, 同时也没有充分考虑低层胶囊对高层胶囊的影响. 文献[11]通过采用注意力路由来调整训练参数的大小进而改变不同空间位置上胶囊的权重, 虽然在一定程度上增加了对重要胶囊的关注, 但没有充分考虑从低层胶囊到高层胶囊的预测过程中低层胶囊的影响. 

由上述分析可见, 尽管胶囊网络是近年来模式识别领域的一大研究热点, 已经取得一些研究成果, 但目前仍处于起步阶段, 有很多尚待完善之处, 例如特征提取不充分、在复杂数据集上的分类效果较差等. 针对以上问题, 本文提出了一种改进的胶囊网络模型, 主要贡献如下: 

1)提出一种多阶段注意力胶囊网络的新模型, 该模型分别在卷积层和动态路由层中引入了注意力机制, 这使得模型的参数可以根据与给定任务相关的图像区域进行更新. 注意力机制考虑了特征之间的相关性, 保证能够学习到更多和任务相关的重要特征, 从而提升了效率. 

2)为充分提取特征信息以及特征之间的空间位置信息, 在卷积层中引入注意力机制. 具体地: 对于高层特征, 重点考虑其包含的高度抽象语义, 因此采用通道注意力(Channel attention, CA)机制; 对于低层特征, 重点考虑特征之间的空间位置信息, 因此采用空间注意力(Spatial attention, SA)机制. 

3)为提高对仿射变换图像的鲁棒性, 提出基于向量的注意力(Vector attention, VA)机制并且应用到胶囊网络动态路由层中的低级胶囊中, 充分考虑初始胶囊(即低级胶囊)之间的相关性, 从而加大对任务相关的初始胶囊的关注, 为高级胶囊的准确预测提供帮助. 

4)传统胶囊网络由于网络架构较浅, 不能充分提取有效特征, 因而在如CIFAR10这样的复杂数据集上效果不好. 针对该问题, 本文提出的多阶段注意力的胶囊网络具有更深的网络架构, 在复杂数据集上也能获得比较满意的结果. 大量的实验结果表明, 改进的胶囊网络模型能够在不同数据集中得到更加准确的分类结果, 明显优于几类常用的胶囊网络模型. 并且, 所提出的胶囊网络在图像重构方面也表现良好. 

图 1  胶囊网络结构图

图 2  多阶段注意力的胶囊网络模型

图 3  CA和SA机制

本文提出的多阶段注意力胶囊网络模型能够有效地解决原始胶囊网络特征提取不充分, 在复杂数据集上表现欠佳的问题. 在特征提取过程中, 我们通过在卷积层中对低层特征采用SA机制, 对高层特征采用CA机制来捕捉有效特征; 在计算效率方面, 我们在动态路由中添加VA机制来更多地考虑和分类任务相关的胶囊; 此外, 胶囊网络能够较好地学习特征间的空间相关性, 从而解决CNN特征间的空间关系难以捕获的问题. 通过实验可以看出, 本文的模型无论在简单数据集还是复杂数据集上都明显优于其他的胶囊网络模型. 未来的工作将专注于更加复杂的数据集以及模型中注意力机制模块的优化, 同时改进图像重构的模型, 得到还原度更高的重构图像, 进而用于模型训练. 

作者简介

宋燕

上海理工大学教授. 2001年获得吉林大学学士学位, 2005年获得电子科技大学硕士学位, 2013年获得上海交通大学博士学位. 主要研究方向为模式识别, 数据分析和预测控制. 本文通信作者. E-mail: sonya@usst.edu.cn

王勇

上海理工大学硕士研究生. 2019年获得皖西学院学士学位. 主要研究方向为图像处理. E-mail: 18856496454@163.com