TrOCR: 基于Transformer的光学字符识别方法，使用预训练模型

温馨提示：
本篇文章已同步至"AI专题精讲" TrOCR: 基于Transformer的光学字符识别方法，使用预训练模型

摘要

文本识别是文档数字化中的一个长期研究问题。现有方法通常基于CNN进行图像理解，基于RNN进行字符级文本生成。此外，通常还需要另一个语言模型作为后处理步骤来提升整体准确率。本文提出了一种端到端的文本识别方法TrOCR，该方法使用预训练的图像Transformer和文本Transformer模型，统一采用Transformer架构进行图像理解和wordpiece级别的文本生成。TrOCR模型结构简单但效果显著，可以通过大规模的合成数据进行预训练，并在人工标注的数据集上进行微调。实验表明，TrOCR模型在印刷体、手写体以及自然场景文本识别任务上都优于当前最先进的模型。TrOCR模型与代码已公开发布，地址为：https://aka.ms/trocr。

1. 引言

光学字符识别（OCR）是将打字、手写或印刷文本图像转换为机器编码文本的电子或机械过程，其数据来源可以是扫描文档、文档照片、场景照片，或图像上叠加的字幕文本。典型的OCR系统包括两个主要模块：文本检测模块和文本识别模块。文本检测旨在定位文本图像中的所有文本块，可以是词级别或行级别。文本检测任务通常被视为目标检测问题，可以使用如YoLOv5和DBNet（Liao et al. 2019）等常规目标检测模型。而文本识别则旨在理解文本图像内容，并将视觉信号转录为自然语言token。文本识别任务通常被建模为一个encoder-decoder问题，其中现有方法采用基于CNN的encoder进行图像理解，采用基于RNN的decoder进行文本生成。本文聚焦于文档图像的文本识别任务，暂不涉及文本检测，作为未来工作方向。

近年来文本识别的研究进展（Diaz et al. 2021）表明，采用Transformer（Vaswani et al. 2017）架构可带来显著性能提升。然而，现有方法仍以CNN作为backbone，并在其之上构建self-attention机制用于图像理解。对于decoder，通常使用CTC（Connectionist Temporal Classification，Graves et al. 2006）结合外部字符级语言模型来提高整体准确率。尽管这种encoder-decoder混合方法取得了不错的成果，但在利用预训练CV和NLP模型方面仍存在提升空间：

1）现有方法中的网络参数都是从头开始在合成或人工标注数据集上训练，未充分利用大规模预训练模型；

2）随着图像Transformer的日益普及（Dosovitskiy et al. 2021; Touvron et al. 2021），特别是近年来的自监督图像预训练方法（Bao, Dong, and Wei 2021），一个自然而然的问题是：是否可以用预训练的图像Transformer取代CNN作为backbone，同时将其与预训练的文本Transformer结合，用于同一个文本识别任务框架中。

为此，我们提出了TrOCR，一种端到端的基于Transformer的OCR模型，用于结合预训练的CV和NLP模型进行文本识别，如图1所示。不同于现有的文本识别模型，TrOCR是一个结构简单但有效的模型，不依赖CNN作为backbone。我们遵循（Dosovitskiy et al. 2021）的做法，首先将输入文本图像调整为384×384大小，然后将图像划分为一系列16×16的patch，这些patch作为图像Transformer的输入。标准的Transformer架构及其self-attention机制被用于encoder和decoder两端，最终以wordpiece为单位生成识别文本。为了有效训练TrOCR模型，其encoder可以用ViT风格的预训练模型初始化（Dosovitskiy et al. 2021; Touvron et al. 2021; Bao, Dong, and Wei 2021），而decoder则可以用BERT风格的预训练模型初始化（Devlin et al. 2019; Liu et al. 2019; Dong et al. 2019; Wang et al. 2020b）。因此，TrOCR具有以下三方面优势：

第一，TrOCR使用了预训练的图像Transformer和文本Transformer模型，利用了大规模无标注数据进行图像理解和语言建模，无需外部语言模型；
第二，TrOCR不依赖任何卷积网络作为backbone，也不引入图像特定的归纳偏置，使模型实现和维护都非常简单；
第三，TrOCR在OCR基准数据集上的实验结果表明，该模型在印刷体、手写体及自然场景文本图像识别方面均能达到当前最优性能，无需复杂的前后处理步骤。

此外，TrOCR也可以轻松扩展至多语言文本识别任务，仅需在decoder端使用多语言预训练模型并扩展词表即可。

本文的贡献总结如下：

1. 本文提出了TrOCR，一种结合预训练CV和NLP模型的端到端Transformer文本识别OCR模型。据我们所知，这是首次在OCR文本识别任务中联合利用预训练图像和文本Transformer。
2. TrOCR基于标准的Transformer encoder-decoder模型架构实现，完全不依赖卷积结构，也不需要复杂的前/后处理步骤，性能达到SOTA水平。
3. TrOCR模型与代码已公开发布，地址为：https://aka.ms/trocr。

2. TrOCR

2.1 模型架构

TrOCR 基于 Transformer 架构构建，包括用于提取视觉特征的 image Transformer 和用于语言建模的 text Transformer。我们在 TrOCR 中采用了标准的 Transformer 编码器-解码器结构。编码器用于获取图像 patch 的表示，解码器则在视觉特征和之前预测的引导下生成 wordpiece 序列。

编码器
编码器接收一个输入图像 $x_{img}\in {\mathbb{R}}^{3\times H_0\times W_0}$，并将其调整为固定大小 $(H,W)$。由于 Transformer 编码器无法直接处理原始图像，除非将其转化为一系列输入 token，因此编码器将输入图像划分为 $N=HW/P^2$ 个大小固定为 $(P,P)$ 的正方形 patch，同时保证调整后的图像宽度 $W$ 和高度 $H$ 能够被 patch 大小 $P$ 整除。随后，这些 patch 会被展平为向量并线性投影为 $D$ 维向量，也就是 patch embedding。$D$ 是贯穿整个 Transformer 各层的隐藏维度大小。

与 ViT（Dosovitskiy et al. 2021）和 DeiT（Touvron et al. 2021）类似，我们保留了通常用于图像分类任务的特殊 token “[CLS]”。“[CLS]” token 汇聚了所有 patch embedding 的信息，代表整个图像。同时，在使用 DeiT 预训练模型初始化编码器时，我们也保留了输入序列中的 distillation token，使模型能够从教师模型中学习。这些 patch embedding 和两个特殊 token 都根据其绝对位置赋予了可学习的一维位置编码。

与 CNN 类网络提取的特征不同，Transformer 模型不具有图像特有的归纳偏置，而是将图像作为一系列 patch 进行处理，这使得模型更容易对整个图像或各个独立 patch 分配不同的注意力。

Decoder
我们在 TrOCR 中使用原始的 Transformer decoder。标准的 Transformer decoder 同样由一组结构相同的层堆叠而成，其结构与 encoder 中的层相似，唯一的区别是 decoder 在多头自注意力和前馈网络之间插入了“encoder-decoder attention”，以便对 encoder 的输出施加不同的注意力。在 encoder-decoder attention 模块中，key 和 value 来自 encoder 的输出，而 query 来自 decoder 的输入。此外，decoder 在自注意力中使用 attention masking，以防止其在训练过程中获取超过预测阶段的信息。由于 decoder 的输出相较于输入会右移一个位置，因此 attention mask 需要确保第 i 个位置的输出只能关注前面的输出，即位置小于 i 的输入。
$$h_i=Proj(Emb(Toke{n}_i))$$
$$\sigma (h_{ij})=\frac{e^{h_{ij}}}{{\sum }_{k=1}^V{e}^{h_{ik}}}forj=1,2,\dots ,V$$

decoder 的 hidden states 通过一个线性层从模型维度投影到词表大小 VV 的维度上，然后通过 softmax 函数计算出词表上的概率分布。我们使用 beam search 获取最终输出。

2.2 模型初始化

encoder 和 decoder 均由在大规模有标签和无标签数据集上预训练的公开模型初始化。

Encoder Initialization
TrOCR 模型中的 encoder 初始化采用了 DeiT（Touvron et al. 2021）和 BEiT（Bao, Dong, and Wei 2021）模型。DeiT 使用 ImageNet（Deng et al. 2009）作为唯一的训练集来训练图像 Transformer。作者尝试了不同的超参数和数据增强方法，以提升模型在数据利用方面的效率。此外，他们将一个强大的图像分类器的知识蒸馏到初始 embedding 中的 distillation token，从而在效果上达到了与 CNN 模型竞争的水平。

参考 Masked Language Model 的预训练任务，BEiT 提出了 Masked Image Modeling 任务用于预训练图像 Transformer。每张图像会被转化为两个视角：图像 patch 和视觉 token。他们使用离散 VAE（Ramesh et al. 2021）的 latent code 将原始图像 token 化为视觉 token，随机遮盖一些图像 patch，并训练模型恢复原始的视觉 token。BEiT 的结构与图像 Transformer 相同，但相比 DeiT 不包含 distillation token。

Decoder Initialization
我们使用 RoBERTa（Liu et al. 2019）和 MiniLM（Wang et al. 2020b）模型来初始化 decoder。总体上，RoBERTa 是对（Devlin et al. 2019）的复现研究，系统评估了许多关键超参数和训练数据规模对性能的影响。基于 BERT，RoBERTa 移除了 next sentence prediction 目标，并动态改变了 Masked Language Model 的掩码模式。

MiniLM 是对大型预训练 Transformer 模型的压缩版本，同时保留了 99% 的性能。与以往使用 soft target 概率或中间表示来指导 student 模型训练的方法不同，MiniLM 通过蒸馏 teacher 模型最后一层 Transformer 的 self-attention 模块进行训练，并引入 teacher assistant 协助蒸馏。

将上述模型加载到 decoder 中时，其结构与目标结构并不完全匹配，因为这些模型本身仅为 Transformer 的 encoder 结构。例如，encoder-decoder attention 层在这些模型中是不存在的。为了解决这个问题，我们通过手动设置参数映射关系，将 RoBERTa 和 MiniLM 模型初始化到 decoder 上，对于缺失的参数则采用随机初始化。

2.3 任务流程

在本工作中，文本识别任务的流程为：给定 textline 图像，模型提取视觉特征，并根据图像和已有上下文预测 wordpiece token。ground truth token 序列以 “[EOS]” token 结尾，表示句子的终止。在训练过程中，我们将该序列整体向后偏移一个位置，并在开头添加 “[BOS]” token 表示生成的开始。偏移后的 ground truth 序列被输入到 decoder 中，其输出通过与原始 ground truth 序列计算交叉熵损失进行监督。在推理阶段，decoder 从 “[BOS]” token 开始，迭代生成输出，每次将新生成的 token 作为下一步的输入。

2.4 预训练

我们在预训练阶段使用文本识别任务，因为该任务能够使模型同时学习视觉特征提取和语言建模的能力。预训练过程分为两个阶段，两个阶段所使用的数据集不同。在第一阶段，我们合成了一个包含数亿条打印 textline 图像的大规模数据集，并在此数据集上对 TrOCR 模型进行预训练。在第二阶段，我们构建了两个相对较小的数据集，分别对应于打印文本和手写文本的下游任务，每个数据集包含数百万条 textline 图像。对于 scene text recognition 任务，我们使用已有的、被广泛采用的合成 scene text 数据集。随后，我们在这些任务特定的数据集上分别进行第二阶段的预训练，所有模型均由第一阶段的模型进行初始化。

2.5 微调

除 scene text recognition 实验外，预训练后的 TrOCR 模型会在下游的文本识别任务上进行微调。TrOCR 模型的输出基于 Byte Pair Encoding (BPE)（Sennrich, Haddow, and Birch 2015）和 SentencePiece（Kudo and Richardson 2018），不依赖于任何任务相关的词表。

2.6 数据增强

我们利用数据增强来提升预训练和微调数据的多样性。对于打印体和手写体数据集，使用六种图像变换方式加上保留原图，总共七种方式，分别为：随机旋转（-10 到 10 度）、高斯模糊、图像膨胀、图像腐蚀、缩放和加下划线。我们对每个样本以相同概率随机决定采用哪一种图像变换。对于 scene text 数据集，按照（Atienza 2021）的方法应用 RandAugment（Cubuk et al. 2020），其增强类型包括反色、弯曲、模糊、加噪、变形、旋转等。

温馨提示：
阅读全文请访问"AI深语解构" TrOCR: 基于Transformer的光学字符识别方法，使用预训练模型