TrOCR:基於Transformer的新一代光學字元識別

微軟技術棧 發表於 2021-10-20

編者按:很長一段時間以來,文字識別一直都是一個重要的關於文件數字化的研究課題。現有的文字識別方法通常採用 CNN 網路進行影像理解,採用 RNN 網路進行字元級別的文字生成。但是該方法需要額外附加語言模型來作為後處理步驟,以提高識別的準確率。
為此,微軟亞洲研究院的研究員們展開了深入研究,提出了首個利用預訓練模型的端到端基於 Transformer 的文字識別 OCR 模型:TrOCR。該模型簡單有效,可以使用大規模合成資料進行預訓練,並且能夠在人工標註的資料上進行微調。實驗證明,TrOCR 在列印資料和手寫資料上均超過了當前最先進的模型。訓練程式碼和模型現已開源。希望感興趣的讀者可以閱讀全文,瞭解 TrOCR 的優勢所在!

光學字元識別(OCR)是將手寫或印刷文字的影像轉換成機器編碼的文字,可應用於掃描文件、照片或疊加在影像上的字幕文字。一般的光學字元識別包含兩個部分:文字檢測和文字識別

  • 文字檢測用於在文字影像中定位文字塊,粒度可以是單詞級別或是文字行級別。目前的解決方案大多是將該任務視為物體檢測問題,並採用瞭如 YoLOv5 和 DBNet 的傳統物體檢測模型。
  • 文字識別致力於理解文字影像並將視覺訊號轉換為自然語言符號。該任務通常使用編碼器-解碼器架構,現有方法採用了基於 CNN 網路的編碼器進行影像理解,以及基於 RNN 網路的解碼器進行文字生成。

在文字識別領域中,Transformer 模型被頻繁採用,其結構的優勢帶來了顯著的效率提升。然而,現有方法仍主要採用 CNN 網路作為主幹網路,並在此基礎上配合自注意力機制來理解文字影像;另外,現有方法還依舊採用 CTC 作為解碼器,配合額外的字元級別的語言模型來提高整體的準確率。這種混合模型雖然取得了巨大的成功,但仍然有很大的提升空間

  1. 現有模型的引數是在合成或人工標註的資料上從零開始訓練的,沒有探索大規模預訓練模型的應用。
  2. 影像 Transformer 模型變得愈發流行,特別是最近提出的自監督影像預訓練。現在應當開始探索預訓練的影像 Transformer 是否可以替代 CNN 主幹網路,以及預訓練的影像 Transformer 是否可以在單一網路中配合預訓練文字 Transformer 進行文字識別任務。

因此,微軟亞洲研究院的研究員們聚焦文字識別任務展開了多項研究,並提出了首個利用預訓練模型的端到端基於 Transformer 的文字識別 OCR 模型:TrOCR,模型結構如圖1。

5a115ce839dda54bcd13d4f35e1365d8.png

圖1:TrOCR 模型結構示意圖

與現有方法不同,TrOCR 簡單、高效,沒有將 CNN 作為主幹網路,而是把輸入的文字影像切分成影像切片,然後輸入到影像 Transformer 中。TrOCR 的編碼器和解碼器則均使用了標準的 Transformer 結構以及自注意力機制,解碼器生成 wordpiece 作為輸入影像的識別文字。為了更有效的訓練 TrOCR 模型,研究員們使用了 ViT 模式的預訓練模型和BERT模式的預訓練模型,來分別初始化編碼器和解碼器。
df63cc8d556c91c8ed16b529e3b3963c.png

論文:https://arxiv.org/abs/2109.10282
程式碼/模型:https://aka.ms/trocr

TrOCR 的優勢有三個方面:

  1. TrOCR 使用預訓練的影像和文字模型,利用大規模無標註資料的優勢來進行影像理解和建模語言模型,不需要額外的語言模型介入。
  2. TrOCR 不需要任何複雜的卷積網路來作為主幹網路,更加易於實現和維護。實驗證明,TrOCR 在列印體和手寫體文字識別任務的基準資料集上均超越了當前最先進的方法,不需要任何複雜的預/後處理步驟。
  3. TrOCR 可以很容易地擴充套件為多語言模型,只需要在解碼器端使用多語種預訓練模型即可。此外,通過簡單調整預訓練模型的引數量配置,使雲/端部署變得極為簡便。

實現方法

模型結構

TrOCR 採用了 Transformer 結構,包括影像 Transformer 和文字 Transformer,分別用於提取視覺特徵和建模語言模型,並且採用了標準的 Transformer 編碼器-解碼器模式。編碼器用於獲取影像切片的特徵;解碼器用於生成 wordpiece 序列,同時關注編碼器的輸出和之前生成的 wordpiece。

對於編碼器,TrOCR 採用了 ViT 模式的模型結構,即改變輸入影像的尺寸,並切片成固定大小的正方形影像塊,以形成模型的輸入序列。模型保留預訓練模型中的特殊標記“[CLS]”代表整張圖片的特徵,對於 DeiT 的預訓練模型,同樣保留了對應的蒸餾token,代表來自於教師模型的蒸餾知識。而對於解碼器,則採用原始的 Transformer 解碼器結構

模型初始化

編碼器和解碼器均應用在經過大規模標註/無標註資料預訓練過的公開模型上,從而進行初始化。編碼器採用 DeiT 和 BEiT 模型進行初始化,而解碼器採用 RoBERTa 模型進行初始化。由於 RoBERTa 的模型結構和標準的 Transformer 解碼器不完全匹配,例如缺乏編碼器-解碼器注意力層,因此研究員們將隨機初始化這些在 RoBERTa 模型中不存在的層。

任務流程

TrOCR 的文字識別任務過程是:給定待檢測的文字行影像,模型提取其視覺特徵,並且給予影像以及已經生成的上文來預測對應的 wordpiece。真實文字以“[EOS]”符號結尾,代表句子的結束。在訓練過程中,研究員們向後旋轉真實文字的 wordpiece 序列,並把“[EOS]”符號挪到首位,輸入到解碼器中,並使用交叉熵損失函式來監督解碼器的輸出。在推斷時,解碼器從“[EOS]”符號開始迭代預測之後的 wordpiece,並把預測出的 wordpiece 作為下一次的輸入。

預訓練

研究員們使用文字識別任務作為預訓練任務,因為它可以使模型同時學習到視覺特徵提取和語言模型兩者的知識。預訓練過程分為兩個階段:第一個階段,研究員們合成了一個包含上億張列印體文字行的影像以及對應文字標註的資料集,並在其上預訓練 TrOCR 模型;第二個階段,研究員們構建了兩個相對較小的資料集,分別對應列印體文字識別任務和手寫體文字識別任務,均包含上百萬的文字行影像,並在列印體資料和手寫體資料上預訓練了兩個獨立的模型,且都由第一階段的預訓練模型初始化。

微調

研究員們在列印體文字識別任務和手寫體文字識別任務上微調了預訓練的 TrOCR 模型。模型的輸出基於 BPE(位元組對編碼)並且不依賴於任何任務相關的詞典。

資料增強

為了增加預訓練資料和微調資料的變化,研究員們使用了資料增強技術,合計七種影像轉換方法(包括保持原有輸入影像不變)。對於每個樣例,研究員們在隨機旋轉、高斯模糊、影像膨脹、影像腐蝕、下采樣、新增下劃線、保持原樣中會機會均等地隨機選取一個方法對影像進行變換。

預訓練資料

為了構建一個大規模高質量的資料集,研究員們從網際網路上公開的 PDF 文件中隨機選取了兩百萬張文件頁面。由於這些 PDF 是數字生成的,因此通過把 PDF 文件轉換為頁面影像,然後提取文字行和裁切影像就可以得到高質量的列印體文字行影像。第一階段的預訓練資料一共包含6億8千萬文字行資料。

對於預訓練的第二階段,研究員們使用5427個手寫字型和 TRDG 開源文字識別資料生成工具,合成了大量的手寫文字行影像,並隨機從維基百科的頁面中選取文字。該階段預訓練的手寫資料集包含合成資料和 IIIT-HWS 資料集,共計1800萬的文字行。此外,研究員們還收集了現實世界中的5萬3千張收據照片,並通過商業 OCR 引擎識別了上面的文字,對照片進行了修正和裁剪。同樣,研究員們也使用 TRDG 合成了100萬列印體的文字行影像,並使用了兩種收據字型和其內建的列印體字型。第二階段預訓練的列印體資料集包含330萬的文字行。表1統計了合成資料的規模。

098ebce2daa90fb44eb0a2b5bbe5661c.png
表1:兩階段預訓練的合成資料規模

預訓練結果

首先,研究員們比較了不同編碼器和解碼器的組合,來尋找最佳的模型設定。研究員們比較了 DeiT,BEiT 和 ResNet50 網路作為編碼器的選項。在比較中,DeiT 和 BEiT 均使用了論文中 base 的模型設定。對於解碼器而言,研究員們則比較使用了 RoBERTa-base 初始化的 base 解碼器和使用 RoBERTa-large 初始化的 large 解碼器。作為對照,研究員們對隨機初始化的模型、CRNN 基線模型以及 Tesseract 開源 OCR 引擎也進行了實驗。

表2給出了組合得到的模型結果。BEiT 編碼器和 RoBERTa-large 解碼器表現出了最好的結果。與此同時,結果表明預訓練模型確實提高了文字識別模型的效能,純 Transformer 模型的效能要優於 CRNN 模型和 Tesseract。根據這個結果,研究員們選取了後續實驗的兩個模型設定:TrOCR-base,包含 334M 個引數,由 BEiT-base 編碼器和 RoBERTa-large 解碼器組成;TrOCR-large,包含558M個引數,由 BEiT-large 編碼器和 RoBERTa-large 解碼器組成。

f53de7145888f681335794ddbe4d1d33.png

表2:在 SROIE 資料集上進行的消融實驗結果

表3給出了 TrOCR 模型和 SROIE 資料集排行榜上當前最先進模型的結果。可以看出,TrOCR 模型憑藉純 Transformer 模型超過了當前最先進模型的效能,同時也證實了其不需要任何複雜的預/後處理步驟。基於 Transformer 的文字識別模型在視覺特徵提取上可以與基於 CNN 的模型有近似效能,在語言模型上可與 RNN 相媲美。

80410870d19a6faf5eb399619aee2ee2.png

表3: 大規模預訓練的 TrOCR 模型在 SROIE 印刷體資料集的實驗結果

表4給出了 TrOCR 模型和 IAM 資料集上現有方法的結果。結果顯示現有方法中 CTC 解碼器和額外的語言模型可以帶來顯著的效果提升。通過與(Bluche and Messina, 2017)比較,TrOCR-large 有著更好的結果,這說明 Transformer 解碼器相比於 CTC 解碼器在文字識別任務中更具競爭力,同時已經具有足夠的建模語言模型的能力,而不需要依賴於額外的語言模型

ba35d9504d224fbb04b942d8aff5f5c2.png

表4: 大規模預訓練的 TrOCR 模型在 IAM 手寫體資料集的實驗結果

TrOCR 使用了來自於影像切片的資訊,獲得了與 CNN 網路相似甚至更好的結果,說明經過預訓練的 Transformer 結構可以勝任提取視覺特徵的任務。從結果來看TrOCR 模型使用純 Transformer 結構超過了所有僅使用合成+IAM 資料的方法,同時在沒有使用額外人工標註資料的情況下,取得了和使用人工標註資料的方法相近的結果。

總結

在本文中,研究員們提出了首個利用預訓練模型的端到端基於 Transformer 的文字識別 OCR 模型:TrOCR。不同於現有方法,TrOCR 不依賴於傳統的 CNN 模型進行影像理解,而是利用了影像 Transformer 作為視覺編碼器,利用文字 Transformer 作為文字編碼器。此外,與基於字元的方法不同,研究員們使用 wordpiece 作為識別輸出的基本單元,節省了在額外的語言模型中額外的計算開銷。實驗證明,在沒有任何後處理步驟,且僅使用一個簡單的編碼器-解碼器模型的情況下,TrOCR 在列印文字和手寫文字識別中均取得了目前最先進的準確率。


這裡有更多微軟官方學習資料和技術文件,掃碼獲取免費版!
內容將會不定期更新哦!
208f6785e4bc3f899ded709a80dff426.jpg