代码解析MixFormer: Mixing Features across Windows and Dimensions
1 前言
最近在阅读CVPR2022的论文MixFormer: Mixing Features across Windows and Dimensions。读完论文后我根据官方的paddlepaddle框架的实现写了pytorch版本的简易版的训练代码，尝试复现。
在这里尝试对MixFormer的代码的运行做一个简单的分析。
代码整理后将公布。
官方实现
https://github.com/chensnathan/PaddleClas/blob/release/2.3/ppcls/arch/backbone/model_zoo/mixformer.py
2 总览
论文中MixFormer的整体框架如下图：
2.1 Convolution Stem
输入的图片Images（H,W,3）首先进入一个Convolution Stem得到指定维度C（即embed_dim）的词向量。
这个stage由几个级联的卷积层组成，如下图（省略了bn，激活层和norm）：
ViT模型中最开始的卷积直接将一个patch的数据(p, p, channel)映射为embedding_dim维度的词向量(1，embedding_dim)。这个卷积的卷积核大小为(p,p)，步长stride=p，通道数为embedding_dim
MixFormer这里的卷积操作，
patch_size = p = 4
首先用了三个3*3的卷积，第一个卷积将输入图像的尺寸下采样两倍（stride=p/2=2），通道数上升到embedding_dim/2。下图是示意图：
后续两个卷积保持尺寸不变，通道数不变。
最后一个卷积是一个（p/2，p/2）即2*2的卷积，将尺寸下采样2倍（stride==p/2=2），通道数上升到embedding_dim。每次卷积计算得到的结果是最后特征图的一个点。这个卷积和ViT的卷积操作是类似的，但是参数不一样。
2.2 BasicLayer
将级联的Mixing Blocks和Stride Conv可以看作一个BasicLayer（最后一个Stage没有下采样卷积）。
2.2.1 Mixing Block
- 首先输入的词向量数据会经过一个norm层（LayerNorm）；
- 然后恢复为特征图形状，并且划分窗口。窗口的划分通过操控tensor的形状完成，如下：
输入是[B, H, W, C]，接下来的变化如下
[B, H/w, w, W/w, w, C] --> [B, H/w, W/w, w, w, C] ---> [B*H/w*W/w, w, w, C]
- 窗口划分后，将一个窗口的特征图展平得到对应的词向量形式，接着在单个窗口上进行自注意力机制计算。self-attention的计算的分析可以见b站霹导的博客（在此膜拜加感谢一波）详解Transformer中Self-Attention以及Multi-Head Attentionhttps://blog.csdn.net/qq_37541097/article/details/117691873
在self-attention的计算过程中会按照下图的设计穿插depth-wise conv的计算以及双向的信息流通。
图中的projection是一个1*1的conv
Mixing Attention的计算流程：
首先是LayerNorm，然后分为两个分支。
self-attention的计算分支中，先通过两个linear映射得到QKV。得到的V与另一个分支传过来的channel interaction相乘，相乘的结果与QK计算得到的自注意力权重相乘。
dw_conv的计算中使用一个linear映射，然后通过depth-wise conv计算，结果可用于产生channel interaction的信息，也会经过一个conv后与另一个分支传过来的spatial interaction的信息相乘，结果与self-attention的结果进行拼接运算。
我把以上文字流程画成了一个比论文图示更加详细的流程图：
- 计算完成后，词向量转窗口，取消窗口划分。取消窗口划分的形状变化如下
输入是[B*H/w*W/w, w, w, C]，接下来的变化
[B， H/w，W/w, w, w, C] ---> [B， H/w, w, W/w, w, C] ---> [B， H, W, C] ---> [B， H * W, C]
- 取消窗口后与最开始的输入连接（加法），完成一个残差连接。
- 之后是带残差的FFN层，公式化表示：x = x + FFN(x)
2.2.2 Stride Conv
这个层在代码中出现在Basic Layer的downsample参数位置，实现为ConvMerging类。
该层负责将特征图的尺寸再次下采样2倍。使用一个kernel为2*2，stride为2的卷积实现。
尺寸计算: output = floor( ( h - 2 ) / 2 ) + 1 = floor( h/2 - 1) + 1 接近于 h/2
在实现时需要先将词向量的形状修改为特征图的形式，经过卷积计算后，再恢复到词向量的形式。
2.3 Projection Layer 和Classification Head
这两个层都使用Linear实现。Projection Layer将通道数继续升高，更大程度地保留通道维度的信息。Classification Head则将通道数映射到指定的类别数目。他们之间用自适应池化将尺寸降低到1。