2022-08-24 文本识别模型-SVTR
本文来自 SVTR论文学习
1. 算法简介
场景文本识别的目的是将自然图像中的文本转换成数字字符序列，这传达了对场景理解至关重要的高级语义信息。由于文本变形、字体、遮挡、杂乱的背景等方面的变化，这项任务具有挑战性。目前用的广泛的文本识别模型通常包含两个模块：（1）提取图像特征的视觉模型；（2）将视觉特征转录为序列特征的序列模型。典型的CRNN算法就是使用了这种方式，如下图：
CRNN网络架构
CRNN这种方式有三个步骤：
- 视觉信息提取。使用VGG或者ResNet提取图像的视觉信息，这时候的特征图还不具备上下文的语义信息
- 序列信息建模。使用双向的LSTM来对视觉信息进行序列建模，使得不同位置的文字序列得以进行联系，获得了具有语义信息的特征序列
- CTC处理。使用线性层（全连接的是通道维度，所以也可以是1×1卷积）对序列特征进行处理，再使用CTC对其进行解码。
这种架构虽然准确，但复杂且LSTM的效率较低，很多移动设备对LSTM的加速效果并不好，所以在实际的应用场景中也存在诸多限制。随着swin transformer在计算机视觉领域大放光彩，swin的这种金字塔结构（像CNN里面的下采样一样）也被引入到文字识别。
本文所学习的SVTR主要有以下贡献：
- 论文证明，在场景文本识别中，单一的视觉语言模型可以获得具有竞争力甚至更高的精度。由于其效率和跨语言的多功能性，它具有广阔的实际应用前景。
- 提出了SVTR，一种文本定制的识别模型。它引入局部和全局混合块，分别提取笔划特征和字符间相关性，并结合多尺度backbone，形成多粒度特征描述。
- 对公共基准的实证研究证明了SVTR的优越性。SVTR-L在识别英汉场景文本方面都取得了最先进的性能。而SVTR-T是有效的，但也是有效的，在一个NVIDIA1080TiGPU中，参数为6.03M，每幅图像文本平均消耗4.5ms。
2. 网络架构
SVTR网络架构图
参考的swin transformer网络结构如下：
swin transformer网络结构
svtr是一个三级逐步下采样的网络（和swin transformer一样，下采样三次），和CNN架构一样，由block + 下采样模块组成。
其block模块和普通的swin中的block模块一致，都是self-attention + mlp。不同的是，SVTR中self-attention的方式和swin的滑动窗口有一定的差异。
SVTR网络中的具体结构如下：
(1) Patch Embedding
SVTR使用两个卷积进行1/4下采样得到token。不同的是，swin是直接使用一个步长为4的4×4卷积进行无重叠的patch embedding，pytorch代码如下所示：
''' swin transformer 的Patch Embedding '''
class PatchEmbed(nn.Module):
def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
super().__init__()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
self.img_size = img_size
self.patch_size = patch_size
self.patches_resolution = patches_resolution
self.num_patches = patches_resolution[0] * patches_resolution[1]
self.in_chans = in_chans
self.embed_dim = embed_dim
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
if norm_layer is not None:
self.norm = norm_layer(embed_dim)
else:
self.norm = None
def forward(self, x):
B, C, H, W = x.shape
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C
if self.norm is not None:
x = self.norm(x)
return x
def flops(self):
Ho, Wo = self.patches_resolution
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
if self.norm is not None:
flops += Ho * Wo * self.embed_dim
return flops
但是，svtr则是使用两个步长为2的3×3卷积进行有重叠的patch embedding（延续的CNN的作风，感受野更大，提取局部信息的表达能力也会比swin的patch embedding要好），paddle代码如下：
class PatchEmbed(nn.Layer):
def __init__(self,
img_size=[32, 100],
in_channels=3,
embed_dim=768,
sub_num=2):
super().__init__()
num_patches = (img_size[1] // (2 ** sub_num)) * \
(img_size[0] // (2 ** sub_num))
self.img_size = img_size
self.num_patches = num_patches
self.embed_dim = embed_dim
self.norm = None
if sub_num == 2:
self.proj = nn.Sequential(
ConvBNLayer(in_channels=in_channels,
out_channels=embed_dim // 2,
kernel_size=3,
stride=2,
padding=1,
act=nn.GELU,
bias_attr=None),
ConvBNLayer(
in_channels=embed_dim // 2,
out_channels=embed_dim,
kernel_size=3,
stride=2,
padding=1,
act=nn.GELU,
bias_attr=None))
def forward(self, x):
B, C, H, W = x.shape
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
x = self.proj(x).flatten(2).transpose((0, 2, 1))
return x
(2) 下采样模块
SVTR延续了paddleocr里面的CRNN下采样结构，3个下采样模块都只对特征图的高度进行下采样，即：
下采样模块
SVTR使用SubSample模块是一个步长为（2，1）核大小为3×3的普通卷积，之所以对高度进行下采样而不对宽度进行下采样，有两个原因：（1）宽维度所包含的文字信息笔记丰富，下采样会造成较多信息的丢失；（2）CTC解码前的Argmax序列输出长度越大，结果越稀疏（包含的空字符越多，CTC解码时连续相同的文字就很大程度上不会被误判掉）
(3) MixBlock结果
由于两个字符可能略有不同，文本识别严重依赖于字符组件级别的特征。然而，现有的研究大多采用特征序列来表示图像文本。每个特征对应于一个切片图像区域，这通常是有噪声的，特别是对于不规则的文本。它不是描述这个角色的最佳方法。而vision transformer引入了二维特征表示，但如何在文本识别的背景下利用这种表示仍然值得研究。
更具体地说，对于embedded组件，作者认为文本识别需要两种特征。（1）第一个是局部组件模式，如类似笔画的特征。它编码了形态特征和特征不同部分之间的相关性；（2）第二种是字符间的依赖性，如不同字符之间或文本和非文本组件之间的相关性。
因此，作者设计了两个混合块（全局Mix和局部Mix），通过使用不同接收场的self-attention来感知上下文的相关性。
(3-1) Global Mixing
Global Mixing评估所有字符组件之间的依赖性。由于文本和非文本是图像中的两个主要元素，这种通用的Mixing可以建立来自不同字符的组件之间的长期依赖关系。此外，它还能削弱非文本成分的影响，同时提高文本成分的重要性。在数学上，对于上一阶段的字符组件CCi−1，它首先被reshape为一个特征序列。当进入Mixing block时，应用layer norm进行标准化，然后进行多头自注意获取其依赖关系。然后，依次应用LN和MLP进行特征融合。与同时引入跳跃连接，形成全局混合块。如下图所示：
Glob Mixing示意图
Global Mixing本质上就是一个简单的自注意力机制，特征图经过线性变换投影到三个空间，然后q矩阵和k矩阵的转置进行矩阵乘法、softmax操作得到attention矩阵，最后和v矩阵进行矩阵乘法得到输出，paddle代码如下所示：
class Attention(nn.Layer):
def __init__(self,
dim,
num_heads=8,
mixer='Global',
HW=[8, 25],
local_k=[7, 11],
qkv_bias=False,
qk_scale=None,
attn_drop=0.,
proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim**-0.5
self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.HW = HW
if HW is not None:
H = HW[0]
W = HW[1]
self.N = H * W
self.C = dim
if mixer == 'Local' and HW is not None:
hk = local_k[0]
wk = local_k[1]
mask = paddle.ones([H * W, H + hk - 1, W + wk - 1], dtype='float32')
for h in range(0, H):
for w in range(0, W):
mask[h * W + w, h:h + hk, w:w + wk] = 0.
mask_paddle = mask[:, hk // 2:H + hk // 2, wk // 2:W + wk //
2].flatten(1)
mask_inf = paddle.full([H * W, H * W], '-inf', dtype='float32')
mask = paddle.where(mask_paddle < 1, mask_paddle, mask_inf)
self.mask = mask.unsqueeze([0, 1])
self.mixer = mixer
def forward(self, x):
if self.HW is not None:
N = self.N
C = self.C
else:
_, N, C = x.shape
qkv = self.qkv(x).reshape((0, N, 3, self.num_heads, C //
self.num_heads)).transpose((2, 0, 3, 1, 4))
q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
attn = (q.matmul(k.transpose((0, 1, 3, 2))))
if self.mixer == 'Local':
attn += self.mask
attn = nn.functional.softmax(attn, axis=-1)
attn = self.attn_drop(attn)
x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((0, N, C))
x = self.proj(x)
x = self.proj_drop(x)
return x
(3-2) Local Mixing
Local Mixing示意图
如图4(b)所示，Local mixing评估了预定义窗口内组件之间的相关性。其目的是对形态特征进行编码，并建立特征内成分之间的关联，从而模拟对特征识别至关重要的笔画样例特征。
与Global Mixing不同，Local mixing考虑的是每个分量都有一个邻域。与卷积类似，混合是以滑动窗口的方式进行。窗口大小根据经验设置为7×11。
与Global Mixing相比，它实现了自我注意机制来捕获局部模式。其实就是Swin transformer里面的那一套，将全局的self-attention转换为局部的self-attention来计算，只不过swin是为了减少计算量，而SVTR是为了获取更多的局部信息。
而且swin是通过reshape这种类似的方式来进行滑窗，并将不同的窗口累加到通道维度上，而SVTR则是直接使用值为0的mask操作。SVTR这种做法和swin相比，计算复杂度还是比较高。
(4) Merging
其实就是下采样操作，和卷积的下采样一样。因为self-attention的计算量和特征图的宽高有关，宽高太大的话，计算复杂度暴涨，所以SVTR对其进行了下采样操作，在低分辨率的特征图上计算可以减少矩阵乘法的计算复杂度。
（5）Combining and Prediction
在最后一个阶段，使用一个Combining进行维度的压缩。即：
if last_stage:
self.avg_pool = nn.AdaptiveAvgPool2D([1, out_char_num])
self.last_conv = nn.Conv2D(
in_channels=embed_dim[2],
out_channels=self.out_channels,
kernel_size=1,
stride=1,
padding=0,
bias_attr=False)
self.hardswish = nn.Hardswish()
self.dropout = nn.Dropout(p=last_drop, mode="downscale_in_infer")
- 首先将高度维度全局池化为1
- 然后是经过全连接层处理
- 字符被进一步压缩为一个特征序列
- 与CRNN中的合并操作相比，SVTR的合并操作可以避免对一维尺寸非常小的token进行卷积，例如对高度为2的token进行卷积。
- 利用组合特征，通过一个简单的并行线性预测来实现识别。具体地说，使用全连接层生成转录序列
- 理想情况下，相同字符被转录成重复的字符，非文本的组件被转录成一个空白符号。就是一个CTC解码模块。
3. PaddleOCR V3中的SVTR模块
总的来看，整个结构其实就是通过self-attention和线性层完成视觉信息和序列信息的编码一步到位。也就是省掉了CRNN里面那个LSTM模块。但是，在PaddleOCR V3里面只使用了两层的self-attention，还是通过mobilenet提取视觉信息，self-attention进行序列信息转换，没有全部使用transformer模块，大概是为了减少计算复杂度。
class EncoderWithSVTR(nn.Layer):
def __init__(
self,
in_channels,
dims=64, # XS
depth=2,
hidden_dims=120,
use_guide=False,
num_heads=8,
qkv_bias=True,
mlp_ratio=2.0,
drop_rate=0.1,
attn_drop_rate=0.1,
drop_path=0.,
qk_scale=None):
super(EncoderWithSVTR, self).__init__()
self.depth = depth
self.use_guide = use_guide
self.conv1 = ConvBNLayer(
in_channels, in_channels // 8, padding=1, act=nn.Swish)
self.conv2 = ConvBNLayer(
in_channels // 8, hidden_dims, kernel_size=1, act=nn.Swish)
self.svtr_block = nn.LayerList([
Block(
dim=hidden_dims,
num_heads=num_heads,
mixer='Global',
HW=None,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop=drop_rate,
act_layer=nn.Swish,
attn_drop=attn_drop_rate,
drop_path=drop_path,
norm_layer='nn.LayerNorm',
epsilon=1e-05,
prenorm=False) for i in range(depth)
])
self.norm = nn.LayerNorm(hidden_dims, epsilon=1e-6)
self.conv3 = ConvBNLayer(
hidden_dims, in_channels, kernel_size=1, act=nn.Swish)
# last conv-nxn, the input is concat of input tensor and conv3 output tensor
self.conv4 = ConvBNLayer(
2 * in_channels, in_channels // 8, padding=1, act=nn.Swish)
self.conv1x1 = ConvBNLayer(
in_channels // 8, dims, kernel_size=1, act=nn.Swish)
self.out_channels = dims
def forward(self, x):
# for use guide
if self.use_guide:
z = x.clone()
z.stop_gradient = True
else:
z = x
# for short cut
h = z
# reduce dim
z = self.conv1(z)
z = self.conv2(z)
# SVTR global block
B, C, H, W = z.shape
z = z.flatten(2).transpose([0, 2, 1])
for blk in self.svtr_block:
z = blk(z)
z = self.norm(
# last stage
z = z.reshape([0, H, W, C]).transpose([0, 3, 1, 2])
z = self.conv3(z)
z = paddle.concat((h, z), axis=1)
z = self.conv1x1(self.conv4(z))
return z
而且paddleocr v3也没有使用论文里面的AdaptiveAvgPool2D，而是使用卷积的方式，很好的解决了out_char_num问题。总的来看，paddleocr v3并未使用完整的SVTR，而是结合了CNN + self-attention，在保留轻量级结构设计的同时去掉了LSTM，替换为可以并行处理的self-attention
