图像分析新方法:“频域学习”(Learning in the Frequency Domain)
就是省略图像压缩/解压缩中计算量最大的步骤,直接利用频域特征来进行图像推理,减少系统中模块之间的数据传输量,从而提升系统性能。输入的数据量更小,深度神经网络在图像分类/分割任务上的精度反而提升了。
论文地址:https://arxiv.org/pdf/2002.12416.pdf
图1 (a)使用RGB图像作为输入的传统的基于CNN的方法的工作流程。
(b)采用DCT系数作为输入的方法的工作流程。CB表示CPU和GPU/加速器之间所需的通信带宽。
3.算法
在本文中,我们提出了一种频域学习的通用方法,包括数据预处理通道和输入数据大小剪枝方法。
图1显示了我们的方法与传统方法的比较。在传统的方法中,高分辨率的RGB图像通常在CPU上进行预处理,然后传输给GPU/AI加速器进行实时推理。由于RGB格式的未压缩图像通常很大,因此CPU和GPU/AI加速器之间的通信带宽要求通常很高。这种通信带宽可能成为系统性能的瓶颈,如图1(a)所示。
为了降低计算成本和通信带宽的要求,将高分辨率的RGB图像向下采样到更小的图像中,往往会导致信息丢失和推理精度下降。在我们的方法中,高分辨率的RGB图像仍然在CPU上进行预处理。然而,它们首先被转换到YCbCr颜色空间,然后转换到频域。
这与最广泛使用的图像压缩标准(如JPEG)一致。所有相同频率的分量被分组到一个信道中。这样,就产生了多个频率通道。如3.2节所示,某些频率信道对推理精度的影响要大于其他信道。因此,我们建议只保留和传输最重要的频率通道到GPU/AI加速器进行推理。与传统方法相比,该方法要求较少的通信带宽,同时具有较高的精度。
我们证明了频率域的输入特征可以应用于所有现有的空间域CNN模型,并且只需进行最小的修改。具体来说,只需要删除输入CNN层并保留剩余的剩余块。将第一残差层作为输入层,修改输入通道的数量以适应DCT系数输入的尺寸。因此,修改后的模型可以保持与原始模型相似的参数数量和计算复杂度。
基于我们的频域模型,我们提出了一种基于学习的信道选择方法来探索给定CNN模型的谱偏,即,哪些频率分量对后续推理任务更有帮助。这一发现促使我们在保持推理准确性的同时,对琐碎的频率成分进行删减,从而大大减少了输入数据的大小,从而降低了域转换的计算复杂度和所需的通信带宽。
图2 用于频域学习的数据预处理管道
3.1 频域数据预处理
数据预处理流程如图2所示。我们遵循空间域的预处理和扩展流,包括图像调整大小、裁剪和翻转(图2中的空间调整大小和裁剪)。然后将图像转换到YCbCr颜色空间,再转换到频域(图2中的DCT变换),再将相同频率的二维DCT系数分组到一个通道中,形成三维DCT立方体(图2中的DCT整形)。正如3.2节中所讨论的,选择有效的频率通道的一个子集(DCT通道选择如图2所示)。所选频道YCbCr色彩空间的连接在一起,形成一个张量(DCT连接如图2所示)。最后,每个频道都是归一化的均值和方差计算训练数据集。
首先考虑亮度通道(Y channel)。假设我们使用图像压缩标准中默认的8x8作为块的尺寸(blocksize)。对于每一个块(block),我们会得到64个DCT的信号,对应了64个不同的频率分量。如果我们原始的图像的尺寸是W x H, 那么我们将会有W/8 x H/8 个DCT信号组成的块。每个块中相同位置的频率分量可以组成一个尺寸为W/8 x H/8的特征图片(feature map),这样我们会产生8x8=64个feature map。同样的对于Cb和Cr通道,我们也可以各自产生64个feature map。总共产生了64x3=192个feature map,这个过程如Figure 10(a)所示。假设W=H=448, 那么现有的基于频域的feature map的尺寸为56x56x192。现在的问题是如何将这些feature map合理的输入到一个已有的DNN网络结构中,使得feature map的尺寸和已有DNN网络的尺寸吻合。
为了便于理解,我们以ResNet-50作为基础的图像分类举例。ResNet-50通常接受的图片输入尺寸为224x224. 在经过一次convolutional layer (stride=2)和pooling之后,此时网络的feature map的尺寸为56x56,和我们产生的频率信号的feature map尺寸吻合。我们可以将192个56x56的频域feature map全部或者部分直接接在ResNet-50的第一个Residue Block之前,从而达到不改变ResNet-50的结构而实现从频域做机器识别的目的。如果我们从192个feature map中选取的64个,则和一个标准的ResNet-50在这一层的feature map个数相同,则网络结构和ResNet-50达到了完全一致。这个过程如下图所示。值得注意的是,由于我们做了8x8的DCT变换,我们实际输入的图片大小为448x448,是标准ResNet-50输入(224x224)的两倍。正因为我们提高了输入图片在空间域的分辨率,我们在后续的实验中可以得到更好的识别精度。
3.2 基于学习的频道选择
图像压缩理论的基础是人眼对于不同色彩空间和频率分量有不同的敏感度,因此我们可以对Cb,Cr通道降采样,也可以对每个DCT信号中高频分量设置更大的量化区间。然而,我们对于这些频率分量对于机器学习的重要性并不了解。如果我们可以忽略掉其中不重要的feature map,就可以不用将这部分数据从图像解码模块传输到图像推理引擎中,从而节省这部分可能成为瓶颈的带宽。
我们现在面临的问题是如何在这192个feature map中做出选择。类比人类视觉系统,在这192个feature map中,我们能否猜想Y通道和低频分量对应的feature map似乎更重要?如果是这样的话,我们如何确定Y通道中应该选择多少个feature map,而CbCr通道有应当选择多少个feature map?在这一部分,我们提出了利用机器学习中添加gate的方法来学习每一个feature map的重要性。在训练中,不仅图像推理的DNN中的weights被训练出来,同时每一个feature map的重要性也被确定。
上图展示了利用gate来选择重要的feature map的方法。原始所有频率分量组成的feature map的尺寸为WxHxC,其中C代表了feature map的个数,正如我们前面介绍的那样,在实验中使用的是C=192. 每个feautre map通过average pooling将会生成一个1x1xC的特征向量,其中每个数值代表了对应feature map。这个特征向量通过一个fully connected layer生成一个1x1xCx2的特征向量对。每一对数字表示这个对应的feature map是否重要,如果index为0的数字比index为1的数字更大,那么这个feature map被认为不重要,整个feature map将会被忽略而不参与后续DNN的计算;反过来说,如果index为1的数字比index为0的数字更大,那么这个feature map被认为重要,将会参与后续DNN的计算。这个操作等效于使用了一个开关(gate)来控制每一个频率信息组成的feature map是否流通到后续的计算中。
具体而言,由于我们使用了argmax函数来选择更大的index进而表示feature map是否被选择参与计算,我们需要一种特殊的方法在训练中将gradient传播到这C个开关的控制网络中。这种方法名为Gumbel-softmax 。
由于频率分量对应的开关被选为通过的数量决定了输入DNN的数据带宽,我们把选择为通过的开关的比例作为loss function中的一项,另一项就是对应机器学习任务中原始的loss。通过最小化loss function来实现机器学习任务精度和输入DNN数据带宽的平衡。
3.3 静态频率通道选择
现在我们拥有了选择重要的feature map的方式,我们有两种方案来减少从图像解码模块到图像推理引擎的数据带宽,这两种方式我们称之为动态(Dynamic)方式和静态(Static)方式。
所谓动态方式,就是每一个频率分量的选择开关由当前输入的图像决定,这种方法可以自适应每一次图像推理(inference)的不同输入。由于选择开关的网络十分简单,我们可以将其放在图像解码模块中。这样从图像解码模块到图像推理引擎之间只需要对被选择的频率分量对应的feature map进行数据传输,可以极大的减少这两个模块之间的带宽需求(see 3rdrow in Figure 12)。
所谓静态方式,就是我们通过训练(training)得到最重要的一些频率分量。在做图像推理(inference)的时候,我们事先就确定只用这些事先确定好的频率分量对应的featuremap,而不会根据不同的图像进行自适应选择。这种静态方式在inference的时候无需选择开关的网络。这种方式不仅可以节省图像解码模块到图像推理引擎的带宽,还可以在编码模块中忽略不重要的频率分量,进而减少图像编码的计算量,延时,以及网络传输的带宽。值得一提的是,通常网络传输的带宽远小于机器内部组件之间的带宽。
4.结果展示
为了演示基于频率分量的机器学习系统和方法,我们选取了两个有代表性的机器学习任务,即图像分类(image classification)和实例分割(instance segmentation)。
4.1 Image Classification
图像分类的任务是对给定的图像进行类别的区分。我们使用ImageNet作为数据集,其中包括大约128万张训练图像和5万张测试图像,总共1000个类别。我们使用了ResNet-50 和MobilenetV2 作为DNN的网络结构。MobilenetV2使用了depth-wise convolution的方式,极大的减少了深度神经网络的计算量和网络的参数量。因此MobilenetV2有很大的潜力作为移动端图像推理引擎。
经过我们的训练,我们得到了一张不同频率分量重要性的Heat map。下图描述了对应192个频率分量的重要性程度。我们可以看出,和我们直觉一致的是,Y通道的重要性高于Cb和Cr通道,同时低频分量的重要性高于高频分量。利用我们提出的gated方法,我们可以通过一次训练就了解该如何分配带宽资源。以下图为例,我们可以使用14个Y通道,5个Cb通道5个Cr通道,共计24个feature map来做图像分类的任务,而不用将最初的192个feature map的数据都从图像解码模块传输到图像推理引擎中。这样我们的传输带宽需求降低为原来的八分之一,而推理的精度(accuracy)反而从标准ResNet-50的75.78%提升至77.196% (see DCT-24 in Table (a))。同样的,在MobileNetV2的实验中,我们通过选取最重要的24个频率分量,得到的识别精度从标准MobileNetV2的71.702%提升至72.364% (see DCT-24 in Table (b) ). 其它的结果,例如选取64,48,12,6个频率分量的精度也展现在Table14中。值得一提的是,根据Figure 13的展示,我们发现低频分量在识别中相对重要,因此我们也尝试了使用heuristic的方式,选取了一个上三角的区域。例如对于DCT-24,我们直觉选取的是Y channel的[0-13],Cb和Cr channel的[0-4]编号的频率分量。这些频率分量和下图中的heatmap略有不同,但是识别的精度几乎没有区别 (See DCT-64T, DCT-48T, and DCT-24T in Table (a) )。这说明我们并不需要对于每个不同的任务都尝试首先获得heatmap,然后严格的按照heatmap来选取频率分量。我们直观上可以选择低频分量(例如这些上三角得到区域),达到和严格按照heatmap选择相同的精度。
3.2 Instance segmentation
实例分割(instance segmentation)结合了物体检测(object detection) 和语义分割(semanticsegmentation)的需求,它的任务是检测出图像中的每个物体对应的像素点,同时将每一个检测出的物体分类。如下图所示,实例分割任务需要检测出物体(例如,人,足球)的boundingbox,还需要在这个bounding box中将属于该物体的像素标注出来。
我们使用了COCO的数据集,其中包含了约10万张训练图像和5千张测试图像。我们使用了Mask RCNN 作为深度神经网络结构。Mask RCNN首先会检测出物体的bounding box,然后在bounding box内部的每个像素做二分,确定其是否属于该物体。
Table 展示了我们在DCT频域做物体识别和实例分割的精度对比。可以看到,从频域做这两个任务,我们可以提升大约0.8%的精度(37.3%到38.1% 以及 34.2%到35.0%)。
5.总结
在本文中,我们提出了一种频域学习方法(使用DCT系数作为输入),并证明了它在分类、检测和分割等多种任务上的通用性和优越性。我们的方法对现有的接收RGB输入的CNN模型的修改很少,因此可以普遍应用于现有的网络训练和推理方法。
结果表明,与传统的空间降采样方法相比,频域学习方法在预处理阶段能更好地保存图像信息,从而提高了图像的精度。我们提出了一种基于学习的动态信道选择方法,并通过实验证明了CNN模型对低频信道比对高频信道更敏感。实验结果表明,在不降低或降低分类、目标检测和实例分割任务的精度的情况下,使用所提出的信道选择方法可以削减高达87.5%的频率信道。