EAST 算法超详细源码解析(三)、模型构建
Date: 2020/05/18
Author: CW
前言:
趁着今日难得有些许闲暇时间,那么就把模型构建这块输出一下吧!
EAST的模型结构还是比较简单的,One-Stage + Anchor-Free,即使从代码实现的角度来看,也甚是直观,简洁明了。
对于特征提取使用的 backbone,原作使用的是 VGG 和 PVA-Net,这里我会给出基于 VGG 实现的代码解析。另外,由于当时在撸这块代码的时候,正好出现了 ResNeSt,于是就把它稍微改了下用到这里了,本文也会给出相应的代码解析。
EAST 模型框架

EAST 模型主要由三部分组成,分别是特征提取、多尺度特征融合以及预测结果生成。

接下来就分别基于这三个部分去解析。
特征提取
不知道你们是否发现,尽管 VGG 已经是很老的模型了,但依旧很多算法会将它用作 backbone,尽管它可能不是 SOTA 的,但是实现起来简单,而且性能也不差,因此还是深受众多炼丹者们的喜爱。这个世界里,往往越原始的东西越稳定,万物在高度进化与发展的过程中,会产生一定程度的“弱化”,这叫“递弱代偿”。不扯远了,下面就看看特征提取的部分吧。

注意以上红框部分,features 属性是用作特征提取的,这里设置为 vgg 的特征提取模块,并不包括最后的全局池化与全连接层,即去掉了头部(犹如迁移学习一般,通常去掉头部 top layer),这是因为我们仅将 VGG 用作特征提取,这样做就可以缩减模型复杂度。
再来看看 backbone 是 ResNeSt 的情况。

这样改的一个理由是保持和 VGG 相似的结构,使其和 VGG 一样都是在经过一定数量的卷积块(bottleneck 大致也可看作是多个卷积块组成)后使用最大池化层进行下采样。另外就是,最大池化能够提取区域内最明显的语义特征,在识别 ground truth 时有一定效果。
这样之后,第一个卷积块即 conv1 保持和原来一样下采样2倍,而后面的每个 bottleneck layer 修改为不进行下采样,下采样2倍均发生在每层 bottleneck layer 后的最大池化层,最后总共下采样32倍,和 EAST 的设计一致。
最后来看看前向反馈的过程:

注意,由于 VGG 均使用最大池化层进行下采样,ResNeSt 中第一次下采样是在第一个卷积块conv1,而我们在每个阶段保存的是最大池化层输出的结果,因此对于 ResNeSt 来说,out 中保存的特征图即为我们所需的4个特征图,它们的尺寸分别是输入图像的1/4、1/8、1/16、1/32,;而对于 VGG,则需要去掉第一个保存下来的特征图,它的尺寸是原图的1/2。
下一部分解析 backbone 的搭建,分别是 VGG 和 ResNeSt,先来看看 VGG。
Backbone
1、VGG


VGG 的特征提取就如以上所示,特征提取器(上一部分的 extractor)仅使用了 VGG 的如上部分,并不会涉及 VGG 整个模型的前向反馈。
2、ResNeSt
这里仅对 ResNeSt 改造的部分进行说明,只有几处小改动,至于其完整结构的代码可以看 CW 之前的那篇文 你的 ResNet 是时候更新了 —— ResNeSt 来也!
改造需要注意两部分,下采样率和通道数。

在每个 bottleneck layer 后分别加一个最大池化层进行下采样,同时注意要取消每个layer 第一个 bottleneck 的下采样(每个 layer 包含多个 bottleneck,self._make_layer 中的参数 stride 影响的是 layer 中第一个 bottleneck,原本的结构中,layer中只有第一个 bottleneck 会进行下采样)。

红框即为改动的部分。

最后,我这里使用的是 resnest50_fast_4s2x40d 结构,2表示将通道划分为2组,40代表组的宽度,与 ResNeXt 中 Cardinality 的概念对应,4代表在划分2组后进一步划分(split)的组数,关于这些概念的更多详情细节可看我在上面提到的那篇文。

特征融合
玩过 FPN 或者 U-Net 等这种金字塔结构网络的朋友们应该都很熟悉这部分了,这里也是那种玩法,首先将 1/32特征图插值上采样2倍,与1/16特征图在通道这个维度 concat 融合,然后经过一个卷积块压缩下通道数,最后再用一个卷积块进一步抽取特征,接下来以此类推继续和1/8、1/4特征图融合...


看了下面前向反馈的过程,应该就清楚以上卷积层的通道数是怎么回事了。


预测结果生成
这部分应该是最简单的了,主要用1x1调整通道数,与所需的预测结果对应。

虽说简单,但在前向反馈的输出中,还是有些细节需注意下,预测输出代表的意义要和标签生成的结果对应起来,这样才能是模型正确学习到东西,否则就被带歪了...

score map 就不多说了,上图注释中已经解释了。对于 loc,也就是 d1~d4,由于它们对应的是输入图像中每点到其所属文本框边界上下左右的距离,那么它们的取值肯定是要在输入图像尺寸范围内的。另外,对于角度,由于在标签生成时,我们对旋转角的搜索范围是 -π/2~π/2,因此这里输出预测结果时也要对应起来。对于标签生成不太清楚的,可以看这个系列的第一篇文 EAST 算法超详细源码解析(一)、数据预处理与标签生成 。
哦,最后有一点在这里提下,在本文开篇展示的那副 EAST 框架图中,在预测结果生成部分,还有个 QUAD geometry,是8个通道的,每个通道代表的是每点与所属文本框的四个顶点之间的横、纵坐标之差,而本文解析的代码中的预测方式称为 RBOX geometry,因其带有旋转角度,所以叫作旋转框 —— Rotated Bounding Box。在实际的实现中,可根据数据集的情况,选择一种方式来实现。