最大化互信息约束的R-NLN高频工件图像分类模型

许知建,欧阳,李毅,李柏林,熊鹰

(西南交通大学 机械工程学院,四川 成都 610031)

0 引言

现代工业生产中,通常采用基于零件的表面图像进行零件分类[1]。在卷积神经网络(convolutional neural network,CNN)[2]提出之前,常用人工设计的图像特征,如尺度不变特征变换SIFT[3]、方向梯度直方图HOG[4]、模板匹配[5]等完成零件的分类。但人为定义的特征,对于不同种类零件的泛化性较差。近年来,随着数据量的爆炸式增加和计算机硬件性能的飞速发展,基于CNN的深度学习模型在图像分类领域取得了重要的地位。但当样本数量较少时,CNN无法将训练集的数据合理泛化到测试集。基于此,研究者提出了小样本学习(few-shot learning,FSL)[6]。而高频元件精密加工工件(简称为高频工件)因其需定制而产量极低,是该问题下典型的工程应用对象。

为使小样本模型能够有效提取图像特征,VILALTA等[7]提出了基于元学习的模型,在识别新类之前,只训练模型提取特征的能力,这类模型被称为特征提取器。为使模型在提取特征后具有分析特征的能力,研究者提出了基于度量的分类算法[8],在特征提取器之后设置一个度量器(measurer),在统一的特征空间中比较支持集(support set)样本特征和查询集(query set)样本特征之间的距离,最后输出类别概率分布以完成分类。

文献[9]的实验结果表明,对于自然图像的分类,传统小样本模型的分类效果次于传统CNN模型。2018年,WANG等[10]提出了非局部网络(non-local network,NLN),通过非局部操作提取图像的全局特征,该网络被广泛应用于小样本目标识别和分类领域。YUE等[11]将NLN应用于细粒度目标识别和视频分类,取得了良好的效果。LIU等[12]提出使用NLN进行图像恢复,将NLN嵌入到循环神经网络(recurrent neural network,RNN),在训练集更小的情况下取得了较好的性能。黄立勤等[13]提出一种NLN模块结合非局部卷积神经网络和图卷积网络,应用于肺动静脉的自动分离。然而,NLN只完成了由全局到局部的全局提取部分,并且高频工件图像还存在类间差异小、全局噪声大的特点,所以要求模型不仅可以提取任务特征,还具备较好的鲁棒性。

因此,本文提出一种基于最大化互信息(mutual information,MI)约束的关联非局部网络(related non-local network,R-NLN)高频工件分类模型。R-NLN将查询(支持)样本的全局特征和局部特征(local feature),通过NLN进行交叉关联以获得适应于查询任务的任务特征;又通过在损失函数中加入互信息熵正则项,最大化约束任务特征和局部特征之间的互信息。本文的主要贡献有:

1)提出了一种基于关联特征的R-NLN特征提取器,完成对工件表面图像自适应任务特征的提取;

2)提出了一种将最大化互信息作为约束条件的损失函数,提高了模型的鲁棒性。

1 模型结构和相关理论

1.1 R-NLN总体结构

图1 R-NLN总体结构

1.2 基于NLN的任务特征提取器ξ2

ξ2采用了NLN作为模型的嵌入单元来提取图像的任务特征。传统卷积核普遍选用3×3或5×5的规格,通过遍历样本上的所有区域以提取该区域的特征。这类卷积核可有效降低样本的维度,减少模型权重的参数规模,从而保证训练的可行性。但是由于自身只能提取局部区域的特征,忽略了远距离像素点之间的关系,所以难以提取全局特征,因此把这样的网络称为局部网络。

与局部网络相反,非局部网络的特点是不受卷积核规格的限制,采用了非局部单元(non-local block)扩大特征参考的区域。非局部单元的特征提取如下式所示:

(1)

式中:向量x是指输入的图像特征;i和j分别是特征上的i位置和j位置;yi是在i位置上的全局特征值;xi是在i位置的原始特征值;xj是在j位置的原始特征值;f(xi,xj)是计算图像上i位置和j位置的关联度的线性函数;g(xj)是将j位置的原始特征值线性映射为全局特征值的函数;Θ(x)作为归一化函数将输入映射到统一的特征空间。

图2 全局特征提取单元

(2)

(3)

图3 任务特征提取单元

1.3 基于余弦相似度的度量器Ω

(4)

式中:s是相似度函数similarity;i代表向量的第i维。

(5)

式中i和j分别表示第i和第j个样本;d为余弦距离函数。

(6)

1.4 基于最大化互信息约束的损失函数

在ξ2提取任务特征的过程中,会受到遍布于高频工件表面的全局噪声干扰,又由于样本数量稀少,模型在训练后期极有可能对噪声信号产生过拟合。但是ξ1提取的局部特征都是低维度特征,受到干扰的幅度较小,所以考虑在损失函数中利用局部特征抗干扰性较好的特点构建相关正则项,以约束模型在训练时参考局部特征和任务特征之间的互信息,从而预防模型对噪声产生过拟合,提高R-NLN模型的鲁棒性。本文采用交叉熵损失函数:

(7)

基于信息熵理论,互信息被定义为:对于两个随机变量X和Y,如果其联合分布为p(x,y),边缘分布为p(x)和p(y),则X和Y的互信息分布I(X;Y)为

(8)

在R-NLN中,需要最大化局部特征向量fl和任务特征向量ft的互信息,由于信息熵是对一个随机变量不确定性的度量,所以使用信息熵来度量互信息的大小。对于一个离散型的随机变量X～p(x),信息熵被定义为

(9)

上式表明信息熵是随机变量X的函数log(1/p(X))的期望。熵的值越大,表示该随机变量的离散程度越大。本文使用互信息熵的大小来描述局部特征和任务特征之间的互有信息量,当局部特征和任务特征之间的重合度较大时,互信息熵LMI越小,反之亦然。结合式(8)和式(9),提出互信息熵正则项lMI:

(10)

最后基于交叉熵损失函数和互信息正则项,提出最大化互信息约束的损失函数Loss:

(11)

式中:λ是用于平衡交叉熵和互信息熵正则项的超参数;N是查询集样本的数量。

通过梯度反向传播更新特征提取器ξ2中卷积核组K2={ki|k∈R1×1×c,i=1,2,3}的权重W:

(12)

W=W-η(∇lc-e+λ∇lMI)

(13)

式中η是控制每一次反向传播中的梯度下降速度的超参数,即学习率。

在交叉熵损失函数中,加入了基于互信息熵的正则化项后,参数的权重变化将会受到互信息的约束,避免了训练过程中受到干扰信号的影响而导致的过拟合,进而提高了R-NLN模型的鲁棒性。

2 实验

2.1 数据集和实验环境

采用miniImageNet数据集和高频元件精密加工工件数据集(简称为高频工件数据集)进行实验。其中,miniImageNet数据集包含100类的自然图像,每类包含600张彩色图像,共有60 000张图像样本。高频工件数据集是科技专项下制作的数据集,包括50类小样本工件图像,每类包含15张彩色图像,共有750张原始图像样本,部分样本如图4所示。

图4 部分高频工件数据集实例

由于高频工件的样本较小,所以对该数据集采用了随机水平/垂直翻转将数据集扩大2倍。采用16∶4∶5的比例设置验证集、训练集、测试集。表1展示了数据集的基本情况。

表1 数据集基本情况

实验软件采用PyTorch深度学习框架,硬件采用Intel© CoreTM I5-9400F CPU@2.90Ghz、NVIDIA Geforce GTX 1080 GPU、32G运行内存的工作站。训练和测试阶段都基于小样本学习下的标准方法N-wayK-shot进行,所有实验均采用Adam优化算法加速模型的训练速率,超参数λ和η初始值分别设置为0.1和0.01。

2.2 分类性能实验

高频工件在表面纹理、特征形状以及颜色上都未表现出较大差异,属于语义差别较小的图像。而自然图像在上述方面的差异较大。为探究R-NLN的性能是否依赖于类间的语义性差别大小,选取自然图像数据集和高频工件图像数据集作对照实验。

对于自然图像分类,使用小样本领域的经典数据集miniImageNet进行分类测试。在训练过程中,从38 400个训练集样本里随机抽取1 000个任务样本,每个任务即为1个mini-batch,并且每10个任务作为Epoch,所以一共有100个Epoch,即训练次数为100次。

对于高频工件图像分类,使用高频工件数据集进行分类测试。在训练过程中,从1 440个训练集样本里随机抽取1 000个任务样本,每个任务即为1个mini-batch。为了完成对照实验,所以选取每10个任务作为1个Epoch,即训练次数同样为100次。

过程中采取跨任务机制训练R-NLN模型,在每个5-way 1-shot(或5-way 5-shot)任务中,除了每个类的1(或5)张图像外,还会从每个类别中随机抽取4张图像作为查询集。即对于一个5-way 1-shot(或5-way 5-shot)任务,每一个训练集中会有5(或25)张支持集图像和25张查询集图像。

采用小样本领域的3类主流算法进行对比实验,包括:基于迁移学习的方法SSMN[14]和PPA[15];基于元学习的方法MN[16]、PN[17]、RN[18]、MAML[19]和AAN[20];基于图卷积网络的方法TPN[21]和GCN[22]。采用平均分类准确率mAp来评价各种算法的性能。在测试过程中,从测试数据集中随机抽取1 000个任务样本,以重复10次的Top-1的平均分类准确率作为最终的评价指标。在训练中为了优化梯度下降过程,每训练1个Epoch则使学习率减小1/2。由于采用了端到端的训练形式,所以过程中也无需微调超参数。实验结果如表2所示。

表2 分类性能实验结果 单位:%

从表2中可以看出,以5-way 5-shot为参考,在自然图像miniImageNet分类中,R-NLN的mAp比迁移学习算法高2.0%,比图卷积算法高8.6%,比其他元学习方法高11.8%。在高频工件图像分类中,R-NLN的mAp比迁移学习算法高2.8%,比图卷积算法高9.7%,比其他元学习算法高12.6%。

因此可知,由于任务特征的提取和互信息约束的共同作用,提高了模型对于自然图像中关键特征的学习能力,在不依赖于类间语义距离的情况下对于复杂语义图像的理解有比较显著的提升。在高频工件图像中,R-NLN在mAp上不仅比同类的元学习算法高6.1%,而且比图卷积算法高9.6%,比迁移学习算法高5.0%。这是因为在损失函数上采用了互信息约束,对样本表面噪声的抗干扰能力做了特定优化,所以模型的鲁棒性有一定的提升。

2.3 消融实验

R-NLN主要通过两点来提高小样本高频工件的分类准确率:1)提取基于NLN的全局特征和局部特征的自适应任务特征;2)通过损失函数的互信息约束减少模型对于元件表面噪声的拟合。为了能够分析任务特征提取和互信息约束对于各自模型的提升,使用控制变量法对模型进行了消融实验,分别验证了以上两点的有效性。

实验以5-way 1-shot(5-shot)为基准,采用16∶4∶5的比例设置验证集、训练集、测试集。在训练过程中,从1 440个训练集样本里随机抽取1 000个任务样本,每个任务即为1个mini-batch,选取每10个任务作为1个Epoch,即训练次数为100次。使用RN和PN作为对照组,进行了4组实验:①只提取任务特征;②只进行互信息约束;③不提取任务特征+不进行互信息约束;④同时进行任务特征提取和互信息约束。实验结果如表3所示。

表3 在高频工件数据集上的消融实验结果 单位:%

从表3的实验结果可以看出:以5-shot的数据为基准,通过比较对照组PN、RN及①和③,可知R-NLN中的任务特征提取模块提高了10.1%的准确率。因为任务特征的本质是基于全局特征和局部特征的关联度而提取的细粒度特征,它对于查询样本具有更好的适应性,在小样本条件下由于支持样本数量少,模型需要尽可能地理解样本,而提取任务特征在这样的情况下能显著提高模型的分类性能。

通过比较对照组PN、RN及②和③,可知,互信息约束在没有进行任务特征提取的条件下,无法显著提高模型的分类准确率。因为互信息约束本身是通过对于局部特征和任务特征的互信息来消除模型对于噪声的拟合,如果不进行任务特征的提取,那么实际上任务特征就退化为全局特征,相当于约束了全局和局部特征之间的互信息,失去了任务特征对于不同任务之间噪声的适应性,所以并不能显著降低模型对于噪声的拟合。

通过对比①和④可知,互信息约束在提取任务特征的前提下提高了4.3%的准确率。原因是最大化互信息约束有效提高了模型的鲁棒性,弥补了全局特征容易拟合噪声的短板。

综上可知,本文提出的“任务特征提取”和“互信息约束”各司其职,一方面提高了模型对于任务特征的提取能力,另一方面通过最大化互信息约束充分利用了局部特征中的优点,提高了模型的鲁棒性,所以R-NLN模型可以获得更高的分类精度。

3 结语

本文针对高频元件精密加工工件样本稀缺难以训练传统CNN模型的问题,提出了一种最大化互信息约束的R-NLN高频工件图像分类模型。R-NLN将查询(支持)样本的全局特征和局部特征通过NLN进行交叉关联以获得适应于查询任务的任务特征,从而提高了模型对支持样本和查询样本之间重点关联部分的提取能力。此外,针对高频工件图像上全局噪声信号较为明显的特点,通过在损失函数中添加互信息熵正则项,以实现在训练过程中最大化任务特征和局部特征之间互信息的约束,完成模型的鲁棒性优化。对miniImageNet和高频工件数据集的实验结果表明,该模型的分类准确率显著高于其他主流的小样本分类模型。