基于transformer的文本識(shí)別方法

很久很久很久之前（大約是兩年前）接觸到了Transformer，seq2seq模型注意力機(jī)制的另一種實(shí)現(xiàn)，完全摒棄了循環(huán)網(wǎng)絡(luò)，使得速度得到極大提升。

所以就忍不住想，把Transformer應(yīng)用到文本識(shí)別上來是否可行呢？這一想法使得我?guī)滋鞄滓共栾埐凰?，夜不能眠，理論上肯定是可以的，因?yàn)門ransformer也是用于seq2seq模型的；鑒于最近終于空閑了，可以放手實(shí)現(xiàn)一下。

關(guān)于Transformer的理論細(xì)節(jié)，本文就不展開了（后面計(jì)劃著...）。 其結(jié)構(gòu)如下圖所示：

而我們主要關(guān)注的是encoder部分：

基于Transformer文本識(shí)別模型的網(wǎng)絡(luò)架構(gòu)

參考文獻(xiàn)[3]把識(shí)別模型網(wǎng)絡(luò)結(jié)構(gòu)劃分成了四部分，如圖：

這種模塊化的結(jié)構(gòu)使得改進(jìn)識(shí)別模型變得十分便捷，只需要在某一個(gè)階段實(shí)現(xiàn)新的功能模塊即可。 為了便于設(shè)計(jì)自己的模型，先簡(jiǎn)單介紹下四個(gè)部分的功能，更詳細(xì)的介紹可以參閱文末的文獻(xiàn)。

1. Trans: 轉(zhuǎn)換階段，就是將輸入圖像使用STN（文中用的是RARE提出的TPS）對(duì)不規(guī)則的文本進(jìn)行變換，實(shí)踐中表明，使用轉(zhuǎn)換對(duì)于性能和魯棒性都有很好的提升
2. Feat： 特征提取階段，常用的Resnet、VGG等提取輸入圖像的特征
3. Seq：序列建模階段，對(duì)于上一步提取的特征，可以reshape變成序列(Batch size, Seq length, Embedding size)，然后再用BiLSTM得到一個(gè)更好的序列，不過會(huì)極大增加計(jì)算開銷
4. Pred： 預(yù)測(cè)階段，對(duì)上面得到的序列進(jìn)行處理，得到模型的輸出。之前介紹的CTC和Attention就是在這個(gè)階段。

由此可見，結(jié)構(gòu)劃分清楚了之后，在我們?cè)O(shè)計(jì)模型的時(shí)候就非常的有幫助了。 上面四步中，通過1,2,3（其中1和3可以不要）將輸入圖像變成了一個(gè)序列，第四步就是典型的seq2seq結(jié)構(gòu)。 所以，本文的改進(jìn)，那就是在Prediction階段使用Transformer替換掉之前常用的CTC或者Attention。 本文的網(wǎng)絡(luò)結(jié)構(gòu)如下(圖很簡(jiǎn)單，就不畫了)：

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

其中BiLSTM可有可無，對(duì)比實(shí)驗(yàn)中可作為變量。

Transformer的編解碼

選用的字符集為10個(gè)數(shù)字、小寫英文字母和三個(gè)標(biāo)點(diǎn)符號(hào)，分別是“'”、“.”、“-”。 這是根據(jù)具體的應(yīng)用場(chǎng)景決定的，因?yàn)槲覀冎饕怯脕碜R(shí)別K12教育所有的英文課本試卷，而這三個(gè)符號(hào)可用于縮寫、省略和連接，比如：

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

編碼

編碼需要額外添加一個(gè)特殊字符，即序列的結(jié)束符號(hào)，使用'eos'來表示，索引默認(rèn)為0，所以，當(dāng)我們將一個(gè)英語單詞映射到對(duì)應(yīng)的索引時(shí)，首先初始化一個(gè)等長(zhǎng)的張量，第一維代表batch size，第二維代表文本索引，初始值為0，然后分別使用文本索引進(jìn)行填充。

解碼

得到了模型的輸出，首先計(jì)算識(shí)別為每個(gè)字符的最大概率索引位置，假設(shè)我們得到一串索引值為[[11,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]，其shape為(1, 15)，代表batch size為1,最大文本長(zhǎng)度為15（實(shí)際中長(zhǎng)度為26），解碼出來得到的是['a','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos','eos']，找到第一個(gè)'eos'出現(xiàn)的位置，只截取該索引前面部分，最終得到的就是“a”。

訓(xùn)練及結(jié)果

借鑒文獻(xiàn)[3]的訓(xùn)練方法，也使用合成數(shù)據(jù)集“ST”和“MJ”來預(yù)訓(xùn)練，沒用其他技巧，比如學(xué)習(xí)率策略、OHEM等。

如預(yù)期，令我們興奮的是，最終的結(jié)果很不錯(cuò)，比文獻(xiàn)[3]所提出的最佳模型略勝一籌。

場(chǎng)外

我們?cè)诰W(wǎng)上一通搜索有沒有使用Transformer用來做文本識(shí)別的文獻(xiàn)或者先例， 最后找到了一個(gè)，不過arxiv上的論文已經(jīng)被撤回了，github上面有開源的項(xiàng)目，發(fā)現(xiàn)跟我們的實(shí)現(xiàn)很不一樣，同時(shí)該項(xiàng)目作者表示他們的代碼計(jì)算損失是有問題的，并且還在做實(shí)驗(yàn)。

所以，我們暫不打算在arxiv上發(fā)表，就放在這里分享一下。 不久之后源碼也會(huì)放在github上，歡迎持續(xù)關(guān)注。

Reference

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