转载：Building A Mobile Text Recognizer via Truncated SVD-based Knowledge Distillation-Guided NAS

最新推荐文章于 2024-09-25 09:46:06 发布

wuling129

最新推荐文章于 2024-09-25 09:46:06 发布

阅读量160

点赞数

分类专栏：图像处理+python+深度学习文章标签：人工智能 OCR

图像处理+python+深度学习专栏收录该内容

120 篇文章 10 订阅

订阅专栏

本文简要介绍BMVC2023录用论文“Building A Mobile Text Recognizer via Truncated SVD-based Knowledge Distillation-Guided NAS”。该工作提出了一种适用于移动端的轻量级文字识别器。该方法将基于截断奇异值分解（TSVD）的知识蒸馏（KD）集成到神经架构搜索（NAS）过程中，并通过引入新颖CNN结构(MCB) 、通道感知搜索(CAS)、LN改进了 NAS 的搜索空间。在手写识别测试集 IAM、ICDAR2013 和 SCUT-HCCDoc以及印刷识别测试集JS-Printed上，搜索出的模型实现了与现有方法相当的精度，但计算成本显著降低，存储空间更小。尤其是在IAM数据上，该方法得到的学生模型在性能上超越教师模型，并在 x86 和 ARM CPU上速度提高了 10.5 倍，尺寸缩小了 8.2 倍。目前已将该方法应用在通用中英文OCR检测和识别上，并开源了相关模型以及移动(IOS/Android)的推理流程和Demo。

一、研究背景

文本在人类获取和保存信息方面发挥着至关重要的作用,一些流行的智能移动应用程序如纸质账单数字化、扫描翻译等，需要部署基于深度学习的系统在计算资源有限的终端设备上。然而，当前大多数提高文本识别准确性的方法会增加模型的复杂度和参数量，使得它们不适合在有限资源的设备上运行。本文提出了一种新的方法来搜索基于CTC的文本识别器的最优结构，包括扩充和优化NAS方法搜索空间，以及在NAS过程中纳入强大教师的监督，并引入了一种基于TSVD的知识提取方法。此外，与之前包括长短期记忆（LSTM）的方法不同，我们在模型中只使用卷积神经网络（CNN）和全连接（FC）层，这对提升网络的并行性和效率有较大的帮助。

二、方法简述

2.1 整体流程

图1展示了所提方法的整体框架，包含NAS Part和Distillation Part两大部分。NAS Part中采用提出的MCB、Channel Aware Search对NAS搜索空间进行改进和扩展，并在FC分类器之前用LN层替换BN层增强归一化特征的表征能力。Distillation Part中提出了TSVD的蒸馏方法，并与NAS的搜索过程进行有效结合。

2.2 网络架构搜索

2.2.1 ProxylessNAS

给定一个超参数化网络 $N$ ，第 $\iota$ 层有 $m_{\iota }$ 个候选的结构 $O_{\iota } = \left \{ O_{i, \iota } \right \}$ ，每个候选结构 $O_{i, \iota }$ 与一个参数 $\alpha_{i, \iota }$ 关联， $\alpha_{i, \iota }$ 代表该结构的重要性。整个搜索过程可以描述成：

其中 $\mathbf{\alpha}$ 是 $\alpha_{i, \iota }$ 的一种集合，W是超参数化网络 $N$ 的权重。 $L_{train}$ 和 $L_{val}$ 分别是训练集和验证集上的损失函数。该搜索过程 $\mathbf{\alpha}$ 和W是交替训练的，固定 $\mathbf{\alpha}$ 训练超参数化网络 $N$ 的W，然后固定W训练 $\mathbf{\alpha}$

，如此反复。训练完毕后，取每层中 $\alpha_{i, \iota }$ 最大对应的结构形成最终的搜索网络。

2.2.2 Mobile Char Block

在基于CTC的识别模型中，通常是先将输入图片通过CNN网络转换成序列特征，然后逐帧解码出文字信息。尽管每帧的感受野可以覆盖输入图像中的多个字符用以捕捉上下文信息，但同时也丢失了单个字符的细节信息，如图2所示。受此启发，我们提出了一个新颖的Mobile Char Block(MCB)模块，用来加强CNN网络提取中心字符细节特征的能力。MCB在Inverted Residual Block结构的基础上将单分支的Depth-wise卷积替换成多分支的Depth-wise卷积求和，如图3所示，多个分支的卷积核宽度从K逐步递减到1，而高度保持一致。此外，在推理阶段，采用重参数化技术将线性的多分支结构合并成单分支结构，减少推理计算量和参数量。

2.2.3 Channel-Aware Search

ProxylessNAS对于搜索不同形式的CNN结构比较友好，但是，如果增加每种结构的通道数搜索，搜索空间将会成几何级别的扩大，这将导致搜索内存增高和搜索时间增长。为了解决这个问题，提出了信道感知搜索机制。该机制不需要在搜索空间中设置大量的搜索块，而只需要设置每层的最大通道和不同比例的掩码。信道感知搜索过程可以公式化为：

其中 $O_{i }$ 是超参数化网络N的第i层输出； $M_{j}$ 是掩膜函数，将 $O_{i }$ 的通道以比例 $r_{j}$ 置为0； $\beta_{j}$ 是 $M_{j}$

的权重。在训练阶段， $\beta_{j}$ 和 $\mathbf{\alpha}$ 一起优化。训练完成后，取最大 $\beta$ 对应的掩膜比例与最大通道数相乘即得到最优的通道数。

2.2.4 Layer Normalization

在基于CTC的模型中，特征序列中超过一半属于空白类别，这可能会影响BN学习的均值和方差的统计。因此，采用LN替换BN对最终的特征序列进行归一化。LN的均值 $\mu$ 和方差 $\sigma$ 计算方式如下：

其中C是序列特征每帧的维度，a是每帧的元素，j是序列的序号。LN使特征序列独立计算均值和方差，可以避免BN带来的缺陷。

2.3 TSVD-based Knowledge Distillation

知识蒸馏可以有效的缩小复杂模型和轻量化模型之间的差距，传统的基于特征的知识蒸馏方法的Loss可以定义为：

其中 $F_{CNN}$ 是教师或学生CNN网络提取的序列特征，reg是回归器，Wreg是回归器的权重。

本文提出了一种TSVD方法，核心思路是对教师的FC层权重进行SVD分解并截断得到回归器reg的权重Wreg。TSVD的具体形式如下：

其中W是教师FC层的权重， $C^{T} and C^{S}$ 分别是教师和学生的特征维度，通过截断值 $C^{S}$ ，可以得到回归器reg的权重V。基于此可以得到TSVD的Loss形式：

三、主要实验结果

整个算法框架在2 NVIDIA 1080TI GPUs上训练，采用ResNet24LN作为蒸馏的Teacher，搜索出的Student模型转换成MNN模型后在Snap-dragon 888 CPU (ARM) 和Intel i7-4790 CPU (x86)平台上进行性能测试。以下是测试和对比的结果。

如表1所示，尽管搜索模型的AR不是最高的，但是在存储和推理速度方面实现了最佳性能，相比Teacher而言，精度损失不大的前提下，在ARM或x86 CPU上的性能均提升了4-10倍。值得注意的是，在IAM数据集上，搜索出的模型比现有的一些方法获得了更高的精度同时存储空间缩小46倍。

表2对提出的方法与现有的轻量级模型进行了比较，包括手动设计的模型和NAS模型。为了确保比较的公平性，将所有模型中的通道数量乘以适当的宽度因子，以实现类似的FLOPs。从结果能看出，搜索出的模型AR度量上表现优秀。

表3展示了在ProxylessNAS的基础上逐步增加MCB、CAS、LN的效果，在ICDAR2013, IAM, SCUT-HCCDoc and JS-Printed 这四个数据集上分别提升 2.16%, 1.39%, 2.53% and 1.56% ，这证明了优化NAS搜索空间的有效性。

如表4所示，使用和不使用搜索的轻量级模型的性能蒸馏与之前的常规训练相当。然而，使用我们的KD搜索模型实现了显著改进。这表明知识蒸馏可以指导NAS算法找到合适的蒸馏架构训练并实现更好的准确性和速度权衡。我们进一步比较了基于TSVD的蒸馏方法与其他KD方法，包括基于关系的和基于特征的方法。结果显示我们的方法获得了最佳性能。

四、小结

本文介绍了一种基于TSVD蒸馏引导NAS的方法来搜索适用于移动应用场景的轻量化的文本识别器，提出了三种方法：新的CNN结构（MCB）、信道感知搜索(CAS)和LN替换优化NAS的搜索空间，此外还引入了截断奇异值分解（TSVD）到NAS以提高师生知识升华学习，形成一个非常快速和紧凑的学生模型。实验证明了该算法的有效性，在IAM、ICDAR2013、SCUT-HCCDoc以及JS-Printed测试集上，本文搜索的模型比以前的文本识别算法快了12−28.9倍，小了6.7−46倍，同时保持有竞争力的识别率。

五、模型开源

OCR是一个应用价值极高的领域，我们进一步延展了论文的工作，将论文中的算法适当优化分别应用到通用中英文OCR检测和识别上，得到5.2M、7.8M的高精度、超轻量级通用OCR检测和识别模型。其中检测算法采用的是DB算法，识别算法与本论文中一致。我们简易的封装了整图OCR识别流程，并与钉钉IOS应用中的“识别图中文字”功能进行端到端的对比，其内部调用了苹果内置的离线端侧OCR能力，具体对比结果如下表：

表5 通用UGC对比数据

表6 通用文档对比数据

轻量化DB检测模型-通用中英文：表5、6分别对比了通用UGC和通用文档上的平均召回准确率和平均误召回率，结果展示我们搜索出来的轻量化模型具备较高的优势。此外，在iPhone12上，通用UGC图片端到端平均耗时400ms。这说明搜索出来的轻量化模型可以运行在一般算力的CPU设备上，可极大地降低OCR的部署成本，具备较强的实用价值。我们开源了这2个轻量化模型，并且开源了移动端的体验Demo，链接如下：

https://www.modelscope.cn/models/damo/cv_proxylessnas_ocr-detection-db-line-level_damo/summary

轻量化CRNN识别模型-通用中英文：

https://modelscope.cn/models/damo/cv_LightweightEdge_ocr-recognitoin-general_damo/summary

移动端Demo：

https://github.com/AlibabaResearch/AdvancedLiterateMachinery/tree/main/OCR/LiteWeightOCR

原文作者:Weifeng Lin, Canyu xie, Dezhi Peng, Jiapeng Wang, Lianwen Jin, Wei Ding, Cong Yao, Mengchao