\u200E
这里有一份轻量级文字识别技术创新大赛优胜团队的修炼秘籍送给你!
发布日期:2021-11-19T11:58:00.000+0000 浏览量:1319次



轻量级文字识别技术创新大赛 是第二届CSIG图像图形技术挑战赛赛题之一,由百度公司承办。本赛题以文字识别为主题,要求参赛选手建立轻量级OCR模型,在兼顾准确率指标与模型大小的同时,重点考察选手的网络结构设计与训练调优能力,进一步推动中文场景文字识别算法与技术的突破。

比赛链接:
https://aistudio.baidu.com/aistudio/competition/detail/75

OCR文本识别任务作为计算机视觉领域的核心问题之一,旨在找出图像中所有的文字信息,进而确定它们的位置和内容。由于文字有不同的外观、形状、尺度和姿态,加上成像时的光照、遮挡等因素干扰,OCR文本识别一直是计算机视觉领域极具挑战性的问题之一,但是在现实世界中又具有庞大的应用基础。

本赛题的主要特点在于:
1)模型轻量级,限制模型大小在10M以内;
2)数据集有限,仅提供10w训练数据集,且不允许使用额外数据集。


作者有话说



我们的方案并不是本赛题参赛队伍中精度最高的,但此篇经验分享的最大优势在于,把我们参赛的心得体会分享给第一次参赛的同学们,以初学者的视角讲述如何思考、分析、解决一个OCR文本识别任务。本经验分享将从拿到赛题后的初步思路构想开始,介绍我们建模过程中每一步的分析和尝试。另外,我们对数据集本身的改动非常小,主要工作集中在模型设计上,希望可以给大家带来一些帮助。


解题思路



就数据集角度而言,本赛题数据集本身更接近常见的场景文本数据集,并具有如下特点:
1)有很多细长文本,用baseline效果不佳;
2)存在特殊字符和特殊形式字体;
3)部分图像前景区分度低,且图像模糊。

数据集的挑战基本涵盖了主流中文场景文本识别任务中存在的大多数困难(虽然有部分错误标注,但飞桨这个数据集做得还是蛮好的),这意味着需要参考学术界方法设计更优的模型。基于此,我们选择了最近很火的Transformer模型。

就模型角度而言,本赛题限制模型大小在10M以内,而1个常用的3*3卷积、通道512的模型,参数量就有2.25M(这里按照MB计算,512*512*3*3/1024/1024=2.25M)。这还只是参数量,如果不使用量化,参数一般以4位浮点数形式存储,那么1个卷积存储大小就接近9M,所以目前很多学术界前沿的场景文本识别方法并不能直接应用于本赛题。最简单且直接的方法就是降低通道数,同时简化整体结构。另外,为保证模型大小始终在10M以内,模型设计过程中的每一步都需要计算模型大小,以确保参数量没有超标。


设计思路





1.整体流程

这里产生了第一个问题,Transformer模型的参数量极大。Transformer模型分为Encoder和Decoder两个模块,我们这里只取了Encoder模块。诚然Decoder模块的类 Attention机制可以为识别增强强大的语义信息,但对于参数量和精度来说,性价比并不高。CTC会是一个简洁高效且性价比更高的选择。所以我们在编码部分选择了Transformer的Encoder,在解码部分选择了CTC。针对数据集中存在的图片模糊、遮挡、多字体等情况,我们选择Backbone去增强这方面的特征提取能力,并在初始阶段选择Mobilenet。




2.Transformer

Transformer的Encoder需要降低通道数,再做一些精简优化。 这里参考SRN [1] ,我们增加了字符阅读顺序这一硬编码,减少多层后对字符位置信息的缺失: Transformer*2_512(512为通道数,*2表示两层)。 此外,我们还尝试了调参操作,并将调完以后的mobilenet+bilstm+ctc作为baseline,这里bilstm通道数为96,mobilenet为large_0.5,参数量为9.8M,所有实验条件均一致。


序列阶段,使用了Transformer的Encoder结构,代码在Transformer Encoder的基础上加入了字符阅读顺序(与SRN论文一致)。由于downsample步骤的缘故,Transformer Encoder只需要80的通道大小即可获得与512通道同样的性能。

1、Transformer Encoder 2层(经过测试,一层参数量0.1M ,模型保存大小约0.5M)
2、其余配置与SRN论文保持一致

代码如下:
class TransformerPosEncoder(nn.Layer):
    def __init__(self,  max_text_length=35, num_heads=8,
                 num_encoder_tus=2, hidden_dims=96,width=80)
:
        super(TransformerPosEncoder, self).__init__()
        self.max_length = max_text_length
        self.num_heads = num_heads
        self.num_encoder_TUs = num_encoder_tus
        self.hidden_dims = hidden_dims
        self.width =width
        # Transformer encoder
        t = 35 #256
        c = 96
        self.wrap_encoder_for_feature = WrapEncoderForFeature(
#... 省略部分配置参数...
            d_key=int(self.hidden_dims / self.num_heads),
            d_value=int(self.hidden_dims / self.num_heads),
            d_inner_hid=self.hidden_dims*2#2,
#... 省略部分配置参数...
            )

    def forward(self, conv_features):
        # feature_dim = conv_features.shape[1]
        feature_dim = self.width
        encoder_word_pos = paddle.to_tensor(np.array(range(0, feature_dim)).reshape(
            (feature_dim, 1)).astype('int64'))
        enc_inputs = [conv_features, encoder_word_pos, None]
        out = self.wrap_encoder_for_feature(enc_inputs)
        return out



3.Downsample Neck

通过实验可以发现,Transformer能给模型带来接近1%的收益(不公平对比),但参数量会增加30M,直接降低通道数也不是一个非常有效的方法,会导致精度下降2.2%。 为此我们设计了一个非常巧妙的neck,并把它命名为Downsample(这么简单粗暴的命名,只是为了和之前的代码注释统一)。
这个Neck可以说是我们方案的核心,正如名字一样,简单粗暴的解决了通道降低带来的精度下降问题。 得益于Downsample,80通道的Transformer获得了比512通道高出2.1%的收益,同时参数量只是前者的0.25倍。 Downsample的设计初衷是模仿Transformer的FFN层,通过先调高通道,再降低的方式,捕获高级语义,并进行柔性压缩。 但通道升高之后卷积参数量仍非常大,为此我们先使用卷积1*1进行通道的上升,接着使用池化层进行图像的下采样,再使用卷积1*1将通道压缩到指定大小。

模型在u-mobilenet阶段,last_conv输出为4*320*640(W,H,C),Downsample阶段流程如下:

1.通过最大池化(2,2)-> (2*160*640)

2.通过最大池化(2,2)->(1*80*640)
3.降维卷积  640-> 80
4.将模型调整成(B*1*80*80)->(B,80,80) (batch, width, channels)

class Im2Seq_downsample(nn.Layer):
    def __init__(self, in_channels, out_channels=80, patch=2*160, conv_name='conv_patch',**kwargs):
        super().__init__()
        self.out_channels = out_channels
        self.patch = patch
        if out_channels == 512:
            self.use_resnet = True
        else:
            self.use_resnet = False
        self.conv1 = ConvBNLayer(
            in_channels=in_channels,
            out_channels=self.out_channels, #make_divisible(scale * cls_ch_squeeze), 
            kernel_size=1,
            stride=1,
            padding=0,
            groups=1,
            if_act=True,
            act='hardswish',
            name=conv_name)
        self.pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
    def forward(self, x):
        x = self.pool(x)
        B, C, H, W = x.shape
        if H != 1 :
            pool2 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
            x = pool2(x)
        x = self.conv1(x)
        # C = 240
        # H = 1
        # W = 80
        assert H == 1
        conv_out = paddle.reshape(x=x, shape=[-1self.out_channels, self.patch])
        conv_out = conv_out.transpose([021])
        #  (NTC)(batch, width, channels)
        out = conv_out
        return out,H*W



4.Umobilenet设计

设计完上述模型,我们打印了badcase进行分析,发现特征信息识别能力并不强,很多模糊、遮挡文字理论上应该拥有被识别的可能性。 此外,调高mobilenet的scale会使精度提升,部分模糊、遮挡文字也可以被识别,但是也会使模型参数量大幅提升。
为此我们设计了一种增强mobilenet里backbone特征的思路: 优化mobilenet。 U-net在图像特征细节的捕获方面非常有优势,因此我们参考U-net设计了一个新的U-mobilenet去提升特征提取能力并且保证参数量符合预期。
表中为最终对比结果,可以看出U-mobilenet可以提升接近2.3%,效果非常可观。 (表中的U-mobilenet通过部分调参,降低了参数量,但并不会因此提升精度,如减少mobilenet中的SE层数等。




通过表格,可以发现模型整体性能的提升主要依靠模型本身的设计。
  • 数据增强方面,我们增加了mixup和cutout方法,以及SRN的尺寸resize;

  • 参数调整方面,我们使用了64*640尺寸,并将max_len调整到35;

  • 数据增强方面,实验证明有1%左右的稳定提升,在不同尺寸上收益并不一致。

进一步的,我们发现尺寸对精度的收益较为明显,大约能达到3-4%的提升,在不同模型和数据增强上表现也不一致,我们理解这些训练技巧之间可以互相弥补,其他的一些细节和训练技巧可以查看我们的项目地址。

AI Studio:
https://aistudio.baidu.com/aistudio/projectdetail/2011330
Github: 
https://github.com/simplify23/Light-STR-Competition-No.5

最终的提交模型在U-mobilenet+Downsample+Transformer的基础上,初始lr=0.01训练了400 epoch以后,又以lr 0.00001 finetune了150 epoch,最终的提交结果为A榜精度80.78% ,B榜精度79%。


总结与完善方向



我们的方案整体上偏向于设计一个更优质的轻量级中文模型,因此在其他训练技巧上做的稍微薄弱,也没有尝试模型压缩方面的工作,整体模型偏简洁。未来,可以考虑增加更多的训练技巧,并且使用模型压缩等技巧(在更大的模型大小上,我们发现模型本身仍有额外收益)。当然,模型本身也有很多值得探究的部分。


对PaddleOCR的想法



做OCR的同学,可能都知道MMlab实验室也开源了MMOCR框架。我们最近也在尝试使用MMOCR。但比赛的那段时间,我们感觉 PaddleOCR 的整体流程很清晰,可扩展性非常强,学习文档也很友好。

如果是初学者,强烈安利使用 PaddleOCR 上手文本领域,提issue的答复速度也很快。飞桨开源框架v2.0版本在函数上和Pytorch非常接近,如果有过渡的同学也不会有很大障碍。

最后希望 PaddleOCR 可以把更多前沿的学术成果迁移到飞桨框架,吸引更多的同学关注。

[1]D. Yu, X. Li, C. Zhang, T. Liu, J. Han, J. Liu, andE. Ding. Towards accurate scene text recognition withsemantic reasoning networks. InProceedings of theIEEE/CVF Conference on Computer Vision and Pat-tern Recognition, pages 12113–12122, 2


关注公众号,获取更多技术内容~