深度学习之图像分类(二十二)-- AS-MLP网络详解

木卯 于 2021-09-30 发布

深度学习之图像分类(二十二)AS-MLP网络详解

在上一讲 MLP-Mixer 最后,我提出了几个问题: MLP-Mixer 是否可以为分割、识别等下游任务提供太大的帮助呢?MLP-Mixer 能不能也增加 patch 内的信息交流(局部性)呢?在本节中,我们学习 AS-MLP,它便解决了上面所述的问题。

img0

1. 前言

AS-MLP 是上海科技大学和腾讯优图实验室共同合作发表的文章,论文为 AS-MLP: AN AXIAL SHIFTED MLP ARCHITECTURE FOR VISION。纯 MLP 网络架构专注于全局的信息交流,却忽略了局部信息的收集,而在计算机视觉中,局部特征非常重要,他们可以提供传统计算机视觉中定义的底层特征(边缘,纹理,拐点等等)。为此,本论文提出 Axial shifted,通过特征图的沿空间轴向的 shift,使得不同空间位置的特征处于同一个 Channel,此后的 Channel Projection 就可以融合不同位置的特征,实现局部性信息交流(其思想和 ShuffleNet 也有一定程度的相似)。AS-MLP是首个迁移到下游任务的 MLP 架构,其作为 backbone 的思想和 Swin 基本一致,通过将 token 进行叠加实现下采样。AS-MLP 在ImageNet 数据集上取得了优于其他 MLP 架构的性能,在 COCO 检测与 ADE20K 分割任务上取得了与 Swin 作为 backbone 相当(略弱)的性能。

img1

2. AS-MLP 网络结构

AS-MLP 和 Swin 类似,其作为 Backbone 的部分如下图所示。在每个阶段中其实 AS-MLP block 被重复了多次。我们先来全局性地讲一下 AS-MLP怎么工作的:

img2

基于 DeiT 和 Swin 等工作的设置,AS-MLP 也提出了 Tiny,Small 和 Base 三种不同的配置,主要涉及 Token 一开始被 Linear Embedding 的维度和每个 Stage 中 AS-MLP block 重复的次数。这样的话不同网络参数量接近,方便进行比较。

img3

3. AS-MLP Block

3.1 Block 结构

在对 AS-MLP 网络整体有个概念后,我们来看看单个 Block 是怎么设计的。粗眼一看,AS-MLP block 和 Transformer block 的整体结构相似,不过将 Multi-Head Self-Attention 部分替换为了 Axial Shift 模块。单个 Axial Shift 模块包含四个全连接层,分别对应于下图中四个 Channel Projection。

img4

AS-MLP Block 实现的伪代码如图所示:

img5

3.2 Axis Shift

3.2.1 感受野分析

通过轴向移动特征信息,AS-MLP 可以得到不同方向的信息流,这有助于捕获局部相关性。该操作使得我们采用纯 MLP 架构即可取得与 CNN 概念相同的感受野,并可以类似卷积核设置 AS-MLP 模块的感受野尺寸以及扩张因子。下图比较了不同架构模块的感受野范围:

img6

3.2.2 并行串行分析

水平和竖直移动处理已经被解耦开来了,那么将水平和竖直移动串行好还是并行好呢?我认为直观上是并行好,因为这样移动时语义是一致的。作者的实验也发现,并行效果会好一点。至于 shift size = 1,dilation = 1 串行更好其实是没有意义的,因为这时候相当于没有移动,只是两个 $1 \times 1$ 卷积,之所以结果略有不同是跑的时候的合理波动吧。所以只用看后三行,并行更好一些。

img7

3.2.3 padding 分析

特征图旋转肯定会有特征移动出特征图范围,那么这时候是通过 zero-padding 还是 cycle-padding 好呢?我认为直观上是 zero-padding 好。我提供个人的两个观点:

综上分析,zero-padding 简单有效,也可理解,实验效果其实也最好。最后作者选择 stride = 5, dilation = 1 进行后续实验。

img8

3.3 额外补充

有人指出,AS-MLP 与百度的那篇 S2-MLP: Spatial-Shift MLP Architecture for Vision (见下图核心模块)真的非常相似,都是采用了垂直、水平移位方式进行空间信息交互,而且还都是上下左右四个方向。可惜 AS-MLP 并未与 S2MLP 进行对比,反而比较晚(指的是见刊arxiv)的 ViP 进行的对比。

img11

这里顺便将 AS-MLP 与 S2MLP 进行对比,参考知乎 AS-MLP:上海科技&腾讯优图开源首个检测与分割领域MLP架构

4. AS-MLP 与下游任务

MLP-Mixer 为什么不能用于下游任务?它对于 H 和 W 是敏感的,因为会对 token 内进行映射。

CNN 为什么能用于下游任务?它对于 H 和 W 是不敏感的,因为只会对 channel 进行变换。

所以为了将 AS-MLP 应用于下游任务,其内部只能使用 Channel Projection。所以在只能使用通道维度 $1 \times 1$ 卷积的基础上,通过把不同位置的特征重排到同一个 channel 是交流空间信息的做法。所以 Axis Shift 操作很自然也不难想到了。

实验结果表明,AS-MLP 作为下游任务的 Backbone 相比于 Swin 略差。

img9

img10

5. 总结

对于 AS-MLP 的一些反思:十字形感受野好吗?十字形感受野操作,它可以更好的提取局部依赖关系吗?不一定!据我了解,暂无理论证明或者过去的传统计算机视觉表明十字形感受野操作可以更好提取视觉底层特征。我们最近的研究确实发现部分计算机底层特征在计算过程中是可以进行正交方向分解的,例如水平方向进行一个操作并在竖直方向进行不同的操作。但是这个感受野应该是类似 Swin 或者传统卷积的稠密型,这样才能得到鲁棒的计算机底层特征。虽然人的视网膜中存在随机性,也可以被视为稠密感受野上增加部分随机掩码,但是这个直接改成稀疏的十字形感受野也不见得好。

AS-MLP 作者最后选择 stride = 5, dilation = 1 进行后续实验。这个真的就能把握长距离依赖吗?对应感受野见 3.2.1 图所示,这其实和卷积中 $5 \times 5$ 感受野的长宽一致,但是元素数量为 9 个,和卷积中 $3 \times 3$ 感受野的感知能力一致,所以在我看来并没有实现特征的长距离依赖把握。

AS-MLP 的旋转能不能也考虑更为复杂的和 Swin 一样地分窗旋转呢?或者想 S2MLP 一样方形感受野呢?但是这样做了,依然没有能够把握长距离依赖,干嘛不直接使用 CNN 呢

img12

Swin 的窗口大,TNT 也可,但是在 MLP 中如何结合图像局部性和长距离依赖依然是值得探讨的点。

6. 代码

代码出处见 此处。 (代码太长,不止一个文件,涉及 Cuda 优化,建议直接看作者开源的工程更好)