基于transformer的文本识别方法-技术圈

很久很久很久之前（大约是两年前）接触到了Transformer，seq2seq模型注意力机制的另一种实现，完全摒弃了循环网络，使得速度得到极大提升。

所以就忍不住想，把Transformer应用到文本识别上来是否可行呢？这一想法使得我几天几夜茶饭不思，夜不能眠，理论上肯定是可以的，因为Transformer也是用于seq2seq模型的；鉴于最近终于空闲了，可以放手实现一下。

关于Transformer的理论细节，本文就不展开了（后面计划着...）。其结构如下图所示：

而我们主要关注的是encoder部分：

基于Transformer文本识别模型的网络架构

参考文献[3]把识别模型网络结构划分成了四部分，如图：

这种模块化的结构使得改进识别模型变得十分便捷，只需要在某一个阶段实现新的功能模块即可。为了便于设计自己的模型，先简单介绍下四个部分的功能，更详细的介绍可以参阅文末的文献。

Trans: 转换阶段，就是将输入图像使用STN（文中用的是RARE提出的TPS）对不规则的文本进行变换，实践中表明，使用转换对于性能和鲁棒性都有很好的提升
Feat：特征提取阶段，常用的Resnet、VGG等提取输入图像的特征
Seq：序列建模阶段，对于上一步提取的特征，可以reshape变成序列(Batch size, Seq length, Embedding size)，然后再用BiLSTM得到一个更好的序列，不过会极大增加计算开销
Pred：预测阶段，对上面得到的序列进行处理，得到模型的输出。之前介绍的CTC和Attention就是在这个阶段。

由此可见，结构划分清楚了之后，在我们设计模型的时候就非常的有帮助了。上面四步中，通过1,2,3（其中1和3可以不要）将输入图像变成了一个序列，第四步就是典型的seq2seq结构。所以，本文的改进，那就是在Prediction阶段使用Transformer替换掉之前常用的CTC或者Attention。本文的网络结构如下(图很简单，就不画了)：

TPS ==》 Resnet50 ==》 BiLSTM(T/F) ==》 Transformer

其中BiLSTM可有可无，对比实验中可作为变量。

Transformer的编解码

选用的字符集为10个数字、小写英文字母和三个标点符号，分别是“'”、“.”、“-”。这是根据具体的应用场景决定的，因为我们主要是用来识别K12教育所有的英文课本试卷，而这三个符号可用于缩写、省略和连接，比如：

it is ==》 it's
et cetera ==》 etc.
role-play
...

编码

编码需要额外添加一个特殊字符，即序列的结束符号，使用'eos'来表示，索引默认为0，所以，当我们将一个英语单词映射到对应的索引时，首先初始化一个等长的张量，第一维代表batch size，第二维代表文本索引，初始值为0，然后分别使用文本索引进行填充。

解码

得到了模型的输出，首先计算识别为每个字符的最大概率索引位置，假设我们得到一串索引值为[[11,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]，其shape为(1, 15)，代表batch size为1,最大文本长度为15（实际中长度为26），解码出来得到的是['a','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos']，找到第一个'eos'出现的位置，只截取该索引前面部分，最终得到的就是“a”。

训练及结果

借鉴文献[3]的训练方法，也使用合成数据集“ST”和“MJ”来预训练，没用其他技巧，比如学习率策略、OHEM等。

如预期，令我们兴奋的是，最终的结果很不错，比文献[3]所提出的最佳模型略胜一筹。

场外

我们在网上一通搜索有没有使用Transformer用来做文本识别的文献或者先例，最后找到了一个，不过arxiv上的论文已经被撤回了，github上面有开源的项目，发现跟我们的实现很不一样，同时该项目作者表示他们的代码计算损失是有问题的，并且还在做实验。

所以，我们暂不打算在arxiv上发表，就放在这里分享一下。不久之后源码也会放在github上，欢迎持续关注。

Reference

Attention Is All You Need
Robust Scene Text Recognition with Automatic Rectification
What Is Wrong With Scene Text Recognition Model Comparisons? Dataset and Model Analysis