Swin Transformer v2

paper:https://arxiv.org/pdf/2111.09883.pdf

Swin Transformer V2: Scaling Up Capacity and Resolution扩展容量和分辨率

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型,现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制,不采用 RNN 的顺序结构,使得模型可以并行化训练,而且能够拥有全局信息。

本文介绍这篇文章是 Swin Transformer 系列的升级版 Swin Transformer v2。Swin Transformer 是屠榜各大CV任务的通用视觉Transformer模型,它在图像分类、目标检测、分割上全面超越 SOTA,在语义分割任务中在 ADE20K 上刷到 53.5 mIoU,超过之前 SOTA 大概 4.5 mIoU!可能是CNN的完美替代方案。除此之外,本文一并介绍 Swin MLP 的代码实现,Swin Transformer 作者们在已有模型的基础上实现了 Swin MLP 模型,证明了 Window-based attention 对于 MLP 模型的有效性。

Swin Transformer Block 有两种,大致结构和 Transformer Block 一致,只是内部 attention 模块分别是 Window-based MSA 和 Shifted Window-based MSA。Window-based MSA 不同于普通的 MSA,它在一个个 window 里面去计算 self-attention,计算量与序列长度 N=hw 成线性关系。Window-based MSA 虽然大幅节约了计算量,但是牺牲了 windows 之间关系的建模,不重合的 Window 之间缺乏信息交流影响了模型的表征能力。Shifted Window-based MSA 就是为了解决这个问题。将下一层 Swin Transformer Block 的 Window 位置进行移动,得到不重合的 patch。

在 Swin Transformer 的基础上,研究人员进一步开发出了用于底层复原任务的 SwinIR

Swin Transformer v2 原理分析:

Swin Transformer 提出了一种针对视觉任务的通用的 Transformer 架构,MSRA 进一步打造了一个包含3 billion 个参数,且允许输入分辨率达到1560×1560的大型 Swin Transformer,称之为 SwinV2。它在多个基准数据集 (包含 ImageNet 分类、COCO 检测、ADE20K 语义分割以及Kinetics-400 动作分类) 上取得新记录,分别是 ImageNet 图像分类84.0% Top-1 accuracy,COCO 目标检测63.1/54.4 box / mask mAP,ADE20K 语义分割59.9mIoU,Kinetics-400视频动作识别86.8% Top-1 accuracy。

Swin Transformer v2 的核心目的是把 Swin Transformer 模型做大,做成类似 BERT large 那样包含 340M 参数的预训练大模型。在 NLP 中,有的预训练的大模型,比如 Megatron-Turing-530B 或者 Switch-Transformer-1.6T,参数量分别达到了530 billion 或者1.6 trillion。

另一方面,视觉大模型的发展却滞后了。 Vision Transformer 的大模型目前也只是达到了1-2 billion 的参数量,且只支持图像识别任务。部分原因是因为在训练和部署方面存在以下困难:

  • 问题1:训练中的不稳定性问题。在大型模型中,跨层激活函数输出的幅值的差异变得更大。激活值是逐层累积的,因此深层的幅值明显大于浅层的幅值。如下图1所示是扩大模型容量时的不稳定问题。 当我们将原来的 Swin Transformer 模型从小模型放大到大模型时,深层的 activation 值急剧增加。最高和最低幅值之间的差异达到了104。当我们进一步扩展到一个巨大的规模 (658M 参数) 时,它不能完成训练,如图2所示。
图1:扩大模型容量时的不稳定问题
图2:使用 Pre-Norm,当进一步扩展到一个巨大的规模 (658M 参数) 时不能完成训练。
  • 问题2:许多下游视觉任务需要高分辨率的图像或窗口,预训练模型时是在低分辨率下进行的,而 fine-tuning 是在高分辨率下进行的。针对分辨率不同的问题传统的做法是把位置编码进行双线性插值 (bi-cubic interpolation),这种做法是次优的。如下图3所示是不同位置编码方式性能的比较,当我们直接在较大的图像分辨率和窗口大小测试预训练的 Imagenet-1k 模型 (分辨率256×256,window siez=8×8) 时,发现精度显着下降。
图3:不同位置编码方式性能的比较
  • 问题3:当图像分辨率较高时,GPU 内存消耗也是一个问题。

为了解决以上几点问题,作者提出了:

方法1:post normalization 技术:解决训练中的不稳定性问题

把 Layer Normalization 层放在 Attention 或者 MLP 的后面。这样每个残差块的输出变化不至于太大,因为主分支和残差分支都是 LN 层的输出,有 LN 归一化作用的限制。如上图1所示,这种做法使得每一层的输出值基本上相差不大。在最大的模型训练中,作者每经过6个 Transformer Block,就在主支路上增加了一层 LN,以进一步稳定训练和输出幅值。

图4:Swin v2 相对于 Swin Transformer 的改进 (红色部分)

方法2:scaled cosine attention 技术:解决训练中的不稳定性问题

原来的 self-attention 计算中,query 和 key 之间的相似性通过 dot-product 来衡量,作者发现这样学习到的 attention map 往往被少数像素对所支配。所以把 dot-product 改成了 cosine 函数,通过它来衡量 query 和 key 之间的相似性。

\[\operatorname{Sim}\left(\mathbf{q}i, \mathbf{k}_j\right)=\cos \left(\mathbf{q}_i, \mathbf{k}_j\right) / \tau+B{i j}\]
式中, \(B_{i j}\) 是下面讲得相对位置编码, \(\tau\) 是可学习参数。余弦函数是 naturally normalized,因 此可以有较温和的注意力值。

方法3:对数连续位置编码技术:解决分辨率变化导致的位置编码维度不一致问题。

  • 该方法可以 更平滑地传递在低分辨率下预先训练好的模型权值,以处理高分辨率的模型权值。
    我们首先复习下 Swin Transformer 的相对位置编码技术。
    \[\operatorname{Attention}(Q, K, V)=\operatorname{SoftMax}\left(Q K^T / \sqrt{d}+B\right) V\]
    式中, \(B \in \mathbb{R}^{M^2 \times M^2}\) 是每个 head 的相对位置偏差项 (relative position bias),\(Q, K, V \in \mathbb{R}^{M^2 \times d}\) 是 window-based attention 的 query, key 和 value。 window 的大小。

作者引入对数空间连续位置偏差 (log-spaced continuous position bias),使相对位置偏差在不同的 window 分辨率之下可以较为平滑地过渡。

方法4:节省 GPU memory 的方法:

1 Zero-Redundancy Optimizer (ZeRO) 技术:

来自论文:Zero: Memory optimizations toward training trillion parameter models

传统的数据并行训练方法 (如 DDP) 会把模型 broadcast 到每个 GPU 里面,这对于大型模型来讲非常不友好,比如参数量为 3,000M=3B 的大模型来讲,若使用 AdamW optimizer,32为的浮点数,就会占用 48G 的 GPU memory。通过使用 ZeRO optimizer, 将模型参数和相应的优化状态划分并分布到多个 GPU 中,从而大大降低了内存消耗。训练时使用 DeepSpeed framework,ZeRO stage-1 option。

2 Activation check-pointing 技术:

来自论文:Training deep nets with sublinear memory cost

Transformer 层中的特征映射也消耗了大量的 GPU 内存,在 image 和 window 分辨率较高的情况下会成为一个瓶颈。这个优化最多可以减少30%的训练速度。

3 Sequential self-attention computation 技术:

在非常大的分辨率下训练大模型时,如分辨率为1535×1536,window size=32×32时,在使用了上述两种优化策略之后,对于常规的 GPU (40GB 的内存)来说,仍然是无法承受的。作者发现在这种情况下,self-attention 模块构成了瓶颈。为了解决这个问题,作者实现了一个 sequential 的 self-attention 计算,而不是使用以前的批处理计算方法。这种优化在前两个阶段应用于各层,并且对整体的训练速度有一定的提升。

在这项工作中,作者还一方面适度放大 ImageNet-22k 数据集5倍,达到7000万张带有噪声标签的图像。 还采用了一种自监督学习的方法来更好地利用这些数据。通过结合这两种策略,作者训练了一个30亿参数的强大的 Swin Transformer 模型刷新了多个基准数据集的指标,并能够将输入分辨率提升至1536×1536 (Nvidia A100-40G GPUs)。此外,作者还分享了一些 SwinV2 的关键实现细节,这些细节导致了 GPU 内存消耗的显着节省,从而使得使用常规 GPU 来训练大型视觉模型成为可能。 作者的目标是在视觉预训练大模型这个方向上激发更多的研究,从而最终缩小视觉模型和语言模型之间的容量差距。

不同 Swin V2 的模型配置:

  • SwinV2-T: C= 96, layer numbers ={2,2,6,2}
  • SwinV2-S: C= 96, layer numbers ={2,2,18,2}
  • SwinV2-B: C= 128, layer numbers ={2,2,18,2}
  • SwinV2-L: C= 192, layer numbers ={2,2,18,2}
  • SwinV2-H: C= 352, layer numbers ={2,2,18,2}
  • SwinV2-G: C= 512, layer numbers ={2,2,42,2}

对于 SwinV2-H 和 SwinV2-G 的模型训练,作者每经过6个 Transformer Block,就在主支路上增加了一层 LN,以进一步稳定训练和输出幅值。

Experiments

模型:SwinV2-G,3B parameters

Image classification

Dataset for Evaluation:ImageNet-1k,ImageNet-1k V2

Dataset for Pre-Training:ImageNet-22K-ext (70M images, 22k classes)

训练策略:分辨率使用192×192,为了节约参数量。2-step 的预训练策略。首先以自监督学习的方式在 ImageNet-22K-ext 数据集上训练 20 epochs,再以有监督学习的方式在这个数据集上训练 30 epochs,SwinV2-G 模型在 ImageNet-1k 上面达到了惊人的90.17%的 Top-1 Accuracy,在 ImageNet-1k V2 上面也达到了惊人的84.00%的 Top-1 Accuracy,超过了历史最佳的83.33%。

图5:Image classification 实验结果

同时,使用 Swin V2 的训练策略以后,Base 模型和 Large 模型的性能也可以进一步提升。比如 SwinV2-B 和 SwinV2-L 在 SwinV1-B 和 SwinV1-L 的基础上分别涨点0.8%和0.4%,原因来自更多的 labelled data (ImageNet-22k-ext, 70M images), 更强的 Regularization,或是自监督学习策略。

Object detection,Instance Segmentation

Dataset for Evaluation:COCO

Dataset for Pre-Training:Object 365 v2

如下图6所示 SwinV2-G 模型与之前在 COCO 目标检测和实例分割任务上取得最佳性能模型进行了比较。SwinV2-G 在 COCO test-dev 上实现了 63.1/54.4 box/max AP,相比于 SoftTeacher (61.3/53.0) 提高了 + 1.8/1.4。

图6:COCO 目标检测和实例分割任务

Semantic segmentation

Dataset for Evaluation:ADE20K

如下图7所示 SwinV2-G 模型与之前在 ADE20K 语义分割基准上的 SOTA 结果进行了比较。Swin-V2-G 在 ADE20K val 集上实现了 59.9 mIoU,相比于 BEiT 的 58.4 高了 1.5。

图7:ADE20k语义分割任务

Video action classification

Dataset for Evaluation:Kinetics-400 (K400)

如下图8所示 SwinV2-G 模型与之前在 Kinetics-400 动作分类基准上的 SOTA 结果进行了比较。可以看到,Video-SwinV2-G 实现了 86.8% 的 top-1 准确率,比之前的 TokenLearner 方法的 85.4% 高出 +1.4%。

图8:K400视频动作分类任务

对比实验:post-norm 和 scaled cosine attention 的作用

如下图9所示,这两种技术均能提高 Swin-T,Swin-S 和 Swin-B 的性能,总体提高分别为 0.2%,0.4% 和 0.5%。说明该技术对大模型更有利。更重要的是,它们能让训练更稳定。对于 Swin-H 和 Swin-G 模型而言,自监督预训练使用原来的 Swin V1 无法收敛,而 Swin V2 模型训练得很好。

图9:post-norm 和 scaled cosine attention 对比实验结果

PointRend –图像细颗粒分割

title
https://arxiv.org/abs/1912.08193

论文地址: https://arxiv.org/abs/1912.08193

gitlab: https://github.com/zsef123/PointRend-PyTorch

存在的问题

在目前的语义分割网络中存在的问题主要有过采样和现采样。

1.过采样( oversample ):对于图片中低频区域( 属于同一个物体 ),没必要使用 太多的采样点,却使用太多采样点造成过采样;

2.欠采样( undersample ) :对于图片中高频区域( 靠近物体边界 ),如果这些区域的采样过于稀疏,导致分割出的边界过于平滑,不大真实

文章要解决的问题是在实例分割任务中边缘不够精细的问题。以MaskRCNN举例,由于计算量和显存的原因,对于每一个ROIAlign之后的proposal我们一般只会upsample到28*28的分辨率输出mask。这对于绝大多数物体显然是不够的。如果想得到像素级别的精度,我们不得不付出更大的计算和存储代价。那有什么办法可以在低代价下仍然得到精细的分割结果呢?其实很重要的一点是往往这些不准确的部分是在物体的边缘,这些边缘其实只占了整个物体中非常小的一部分。所以基于这样的一个想法,作者提出可以每次在预测出来的mask中只选择Top N最不确定的位置进行细分预测。每个细分点的特征可以通过Bilinear插值得到,每个位置上的classifier通过一个简单的MLP来实现。这其实是等价于用一个1*1的conv来预测,但是对于中心很确定的点并不计算。整体的示意图如下:

PointRend 解决了什么问题?

这篇论文讲了一个很好听的故事,即:把语义分割以及实例分割问题(统称图像分割问题)当做一个渲染问题来解决。故事虽然这么讲,但本质上这篇论文其实是一个新型上采样方法,针对物体边缘的图像分割进行优化,使其在难以分割的物体边缘部分有更好的表现

作为一个小白,那么问题来了:

1、什么是渲染?

2、为什么要把图像分割问题当做渲染问题呢?

要想知道什么是渲染,可以参考:

计算机中所说的「渲染」是什么意思?

简单来说,渲染就是“绘制”,把3D的物体在2D平面上绘制出来。

为什么要把图像分割问题和渲染问题扯在一起呢?因为讲故事好听啊,论文好写嘛….咳咳…不不,是因为二者有类似的问题要解决:即物体边缘难以处理。

具体来说,在图像渲染中,对于多个3D物体,在边缘要判断对于镜头而言谁先谁后,而且还得抗锯齿;而对于图像分割问题,边缘恢复也一直是个麻烦事儿,因为在典型的语义分割网络中(如FCN、DeepLab),在CNN内部一般都会相对输入图像降采样16倍,然后再想办法上采样回去。更细致地说,对于 DeepLabV3+,模型最后直接是一个4倍的双线性插值上采样,这显然对物体边缘的预测十分不利。虽然 DeepLabV3+当时在2017年就达到了秒天秒地的 89%mIoU on VOC2012 test (使用了300M JFT 数据集预训练),至今无人超越(因为JFT 数据集 Google没有公开 \手动滑稽),但显然这个上采样过程仍然存在较大的提升空间。

参考链接:Uno Whoiam:DeepLab 语义分割模型 v1、v2、v3、v3+ 概要(附 Pytorch 实现)

而在实例分割网络中,Mask R-CNN 这货生成的 Mask 才 28×28,要是把这样的 mask 拉伸到 不说多了比如 256×256,还指望它可以很好地预测边缘?我只能说这是在想Peach。

事实上,在图像分割任务上边缘预测不理想这个情况其实在许多前人的工作中都有提及,比如 Not All Pixels Are Equal: Difficulty-Aware Semantic Segmentation via Deep Layer Cascade 中就详细统计了语义分割中,模型最容易误判的 pixel基本上都在物体边缘(如下图右上红色部分标记) 。

而关于上采样其实也有一些前人的工作,如 Decoders Matter for Semantic Segmentation: Data-Dependent Decoding Enables Flexible Feature Aggregation,在实现上有点像超分辨率网络 ESPCN 里使用的 sub-pixel convolutional layer 的操作,不过多加了一个二阶范数约束:

总的来说,图像分割边缘预测是一个未被很好解决的问题,而何恺明团队的 PointRend 是对此问题的一个新的思路和解法,接下来将介绍 PointRend 是如何 work 的。

文主要贡献

1.提出可嵌入主流网络的PointRend模块,提高了图像分割精度。

2.把图像分割问题看作渲染问题,本质上是一个新型上采样方法,为图像分割提供独特视角。

3.降低了训练所需的算力。输出224×224分辨率图像,PointRend只需0.9B FLOPs。

二、总体思路

PointRend 方法要点总结来说是一个迭代上采样的过程:

while 输出的分辨率 < 图片分辨率:

  1. 对输出结果进行2倍双线性插值上采样得到 coarse prediction_i。(粗分辨率预测)
  2. 挑选出 N 个“难点”,即结果很有可能和周围点不一样的点(例如物体边缘)。
  3. 对于每个难点,获取其“表征向量”,“表征向量”由两个部分组成,其一是低层特征(fine-grained features),通过使用点的坐标,在低层的特征图上进行双线性插值获得(类似 RoI Align),其二是高层特征(coarse prediction),由步骤 1 获得。
  4. 使用 MLP 对“表征向量”计算得到新的预测,更新 coarse prediction_i 得到 coarse prediction_i+1。这个 MLP 其实可以看做一个只对“难点”的“表征向量”进行运算的由多个 conv1x1 组成的小网络。

整个过程可以这么理解:

小明同学做题,现在有已知条件(coarse prediction_0,fine-grained features),想求解答案(coarse prediction_k),发现直接求(双线性插值or其它方法)不够准确,那就一步一步来吧(从coarse prediction_1,coarse prediction_2….求到coarse prediction_k)。好的,现在求coarse prediction_1,诶,发现有好多东西不知道,不能从 coarse prediction_0 直接得到怎么办?那就找出不知道的(“难点”),在 fine-grained features 里面找出对应的线索(ROIAlign-like 双线性插值),然后在结合 coarse prediction_0 得到整体线索(“特征向量”)求解(使用MLP计算),嗯,终于得到 coarse prediction_1了。再用同样的思路反复求解,直到 coarse prediction_k。

示意图如下:

对于一个coarse prediction(4×4大小),将其上采样两倍(8×8大小,这里可以理解为检测头的输出)后,取了一些难分割的点(大多是边缘部分),取这些点的特征向量输入到MLP网络中,进行point prediction,得到每一个点的新类别,最后结果输出(8×8大小,边缘更加精确的结果)。

另外,其PointRend 训练为了节省时间,没有使用上述的迭代过程,而是使用多种组合的采样方法,不赘述,详见paper。

  1. 从PointRend的应用思路中可以看到,这里包含了两个阶段的特征处理,分别是fine-grained features和coarse prediction部分,如果主干网络是ResNet,那么fine-grained features就是ResNet的stage2输出,也就是4倍下采样时的精细分割结果,而coarse prediction就是检测头的预测结果(还未上采样还原成原图的结果)。
  2. 从coarse prediction中挑选N个“难点”,也就是结果很有可能和周围点不一样的点(比如物体边缘的点)。对于每一个难点,获取他的“特征向量”,对于点特征向量(point features),主要由两部分组成,分别是fine-grained features的对应点和coarse prediction的对应点的特征向量,将这个两个特征向量拼接成一个向量。
  3. 接着,通过一个MLP网络对这个“特征向量”进行预测,更新coarse prediction。也就相当于对这个难点进行新的预测,对他进行分类。

看完这个,我们就可以这么理解,将预测难的点(边缘点)提取出来,再提取其特征向量,经过MLP网络,将这个点的归属进行分类,然后提升这些点的分类准确率。这就是PointRend的思想。

一个PointRend模块包括三部分

1.point selection strategy:用于inference和traing的点选择

对于点采样过程,需要对模型的Train过程和Inference过程做区分

该方法的核心思想是灵活自适应地选择图像平面上的点来预测分割标签。直观地说,这些点应该更密集地位于高频区域附近,例如物体边界,类似于射线追踪中的反混叠问题。我们产生了推理训练的想法。

  • inference推理

通过仅在与其邻域有显着不同的位置进行计算,该方法可用于有效地渲染高分辨率图像(例如,通过光线跟踪)。对于所有其他位置,通过对已经计算的输出值(从粗网格开始)进行插值来获得值。

对于每个区域,我们以粗到精的方式迭代地“渲染”输出蒙版。在规则网格上的点上进行最粗糙级别的预测(例如,通过使用标准的粗糙分段预测头)。在每次迭代中,PointRend使用双线性插值对其先前预测的分段进行上采样,然后在此较密集的网格上选择N个最不确定的点(例如,对于二进制掩码,概率最接近0.5的那些)。然后,PointRend为这N个点中的每一个点计算特征,并预测它们的标签。重复该过程,直到将分段上采样到所需的分辨率为止。

  • training

对于Train过程的点采样操作,同样可以遵循Inference中的操作。但是作者发现,这样子采样对于梯度的传播不太友好,于是只能被迫选择其他的点采样策略——干脆就用随机采样的方式来进行采样。

在训练过程中,PointRend还需要选择一些点,以在这些点上构建用于训练point head的逐点(point-wise)特征。原则上,点选择策略可以类似于推理inference中使用的细分策略。但是,细分引入了一系列步骤,这些步骤对于通过反向传播训练神经网络不太友好。取而代之的是,为了训练,我们使用基于随机采样的非迭代策略

采样策略在特征图上选择N个点进行训练。它旨在使用三个原理将选择偏向不确定区域,同时还保留一定程度的均匀覆盖。对于训练和推理选择,N的值可以不同。

(i)过度生成:我们通过从均匀分布中随机采样kN个点(k> 1)来过度生成候选点。(ii)重要抽样:通过对所有kN个点的粗略预测值进行插值并计算任务特定的不确定性估计,我们将重点放在具有粗略预测的点上。从kN个候选中选择最不确定的βN个点(β∈[0,1])。(iii)覆盖范围:从均匀分布中采样剩余的(1-β)N点。我们用不同的设置来说明此过程,并将其与常规的网格选择进行比较,如下图所示。

在训练时,预测和损失函数仅在N个采样点上计算(除粗略分割外),这比通过细分步骤进行反向传播更简单,更有效。这种设计类似于在Faster R-CNN系统中对RPN + Fast R-CNN的并行训练,其推理是顺序的。

2. Point-wise Representation:逐点表示

PointRend通过组合(例如,级联)两种特征类型(细粒度和粗略预测特征)在选定点上构造逐点特征,如下所述。

  • 细粒度特征

为了允许PointRend呈现精细的分割细节,我们从CNN特征图中提取每个采样点的特征向量。 因为一个点是“实值2D坐标”,所以我们按照标准做法对特征图执行双线性插值,以计算特征向量。 可以从单个特征图中提取特征(例如,ResNet中的res2);也可以按照Hypercolumn方法,从多个特征图(例如res2到res5)中提取并连接它们。

  • 粗预测特征

细粒度的特征可以解析细节,但在两个方面也有不足:

首先,它们不包含特定区域的信息,因此,两个实例的边界框重叠的相同点将具有相同的细粒度特征。但是,该点只能位于一个实例之中。 因此,对于实例分割的任务,其中不同的区域可能针对同一点预测不同的标签,因此需要其他区域特定的信息。

其次,取决于用于细粒度特征的特征图,这些特征可能只包含相对较低级别的信息(例如,我们将对res2使用DeepLabV3)。 因此,需要有更多具有上下文和语义信息的特征。

基于这两点考虑,第二种特征类型是来自网络的粗分割预测,例如表示k类预测的区域(box)中每个点的k维向量。通过设计,粗分辨率能够提了更加全局的上下文信息,而通道则传递语义类别。这些粗略的预测与现有架构的输出相似,并且在训练过程中以与现有模型相同的方式进行监督。例如,在mask R-CNN中,粗预测可以是一个轻量级的7×7分辨率Mask头的输出。

点特征向量(point features),主要由两部分组成,分别是fine-grained features的对应点和coarse prediction的对应点的特征向量,将这个两个特征向量拼接成一个向量

3. point head

给定每个选定点的逐点特征表示,PointRend使用简单的多层感知器(MLP)进行逐点分割预测。这个MLP在所有点(和所有区域)上共享权重,类似于图卷积或PointNet。由于MLP会预测每个点的分割标签,因此可以通过特定任务的分割loss进行训练。

三、效果如何?

实验结果

  • 网络设计

实验使用ResNet-50+ FPN 的Mask-Rcnn作backbone。 Mask-RCNN中的默认head是region-wise FCN,用“ 4×conv”表示,作为用来与本文网络进行比较的基准网络。

为了计算粗略预测,我们用重量更轻的设计替换4×conv Mask头,该设计类似于Mask R-CNN的box head产生7×7Mask预测。具体来说,对于每个边界框,我们使用双线性插值从FPN的P2层提取14×14特征图。这些特征是在边界框内的规则网格上计算的(此操作可以看作是RoIAlign的简单版本)。接下来,我们使用具有256个输出通道步幅为2的 2×2卷积层,后跟ReLU, 将空间大小减小到7×7。最后,类似于Mask R-CNN的box head,用两个带1024宽的隐藏层的MLP为K类分别产生7×7的Mask预测。ReLU用于MLP的隐藏层,并且Sigmoid激活函数应用于输出。

PointRend:在每个选定点上,使用双线性插值从粗预测头的输出中提取K维特征向量,PointRend还从FPN的P2级别插值256维特征向量,步长为4。这些粗预测和细粒度特征向量是串联在一起的,我们使用具有256个通道的3个隐藏层的MLP在选定点进行K类别预测。在MLP的每个层中,我们用K个粗预测特征补充到256个输出通道中,作为下一层输入向量。在MLP中使用ReLU,并将Sigmoid激活函数应用于输出。

不得不说这个针对物体边缘进行优化的上采样方法的确在感官上和数据上都有很不错的效果:

语义分割结果:

实例分割结果(基于MaskR-CNN):

PointRend的一些代码和实现

摘自: https://chowdera.com/2022/194/202207120607167479.html

代码详解: https://www.361shipin.com/blog/1536592971120508928

  • 作者提出可以在预测出来的mask中只选择Top N最不确定的位置进行细分预测。

具体为先根据粗糙预测出来的mask,将mask按类别预测分数排序,选出分数高的前2 类别的mask,计算出在2个类别mask上均有较高得分的Top K个像素点作为K 个不确定点【1个像素点只能对应1个类别,如果它对应2个类别的分数都很高,说明它很可能是边界点,也是不确定的】

def sampling_points(mask, N, k=3, beta=0.75, training=True):
    """
    主要思想:根据粗糙的预测结果,找出不确定的像素点
    :param mask: 粗糙的预测结果(out)   eg.[2, 19, 48, 48]
    :param N: 不确定点个数(train:N = 图片的尺寸/16, test: N = 8096)    eg. N=48
    :param k: 超参
    :param beta: 超参
    :param training:
    :return: 不确定点的位置坐标  eg.[2, 48, 2]
    """
    assert mask.dim() == 4, "Dim must be N(Batch)CHW"   #this mask is out(coarse)
    device = mask.device
    B, _, H, W = mask.shape   #first: mask[1, 19, 48, 48]
    mask, _ = mask.sort(1, descending=True) #_ : [1, 19, 48, 48],按照每一类的总体得分排序
    if not training:
        H_step, W_step = 1 / H, 1 / W
        N = min(H * W, N)
        uncertainty_map = -1 * (mask[:, 0] - mask[:, 1])
        #mask[:, 0]表示每个像素最有可能的分类,mask[:, 1]表示每个像素次有可能的分类,当一个像素
        #即是最有可能的又是次有可能的,则证明它不好预测,对应的uncertainty_map就相对较大
        _, idx = uncertainty_map.view(B, -1).topk(N, dim=1) #id选出最不好预测的N个点
        points = torch.zeros(B, N, 2, dtype=torch.float, device=device)
        points[:, :, 0] = W_step / 2.0 + (idx  % W).to(torch.float) * W_step    #点的横坐标
        points[:, :, 1] = H_step / 2.0 + (idx // W).to(torch.float) * H_step    #点的纵坐标
        return idx, points  #idx:48 || points:[1, 48, 2]
  • 得到不确定点的位置以后,可以通过Bilinear插值得到对应的特征,对每个不确定点的使用一个MLP来进行单独进行细分预测【训练与预测有所区别】。

具体为:通过刚刚得到的不确定点所在图片的相对位置坐标来找到对应的特征点,将此点对应的特征向量与此点的粗糙预测结果合并,然后通过一个MLP进行细分预测。

##训练阶段
def forward(self, x, res2, out):
        """
        主要思路:
        通过 out(粗糙预测)计算出top N 个不稳定的像素点,针对每个不稳定像素点得到在res2(fine)
        和out(coarse)中对应的特征,组合N个不稳定像素点对应的fine和coarse得到rend,
        再通过mlp得到更准确的预测
        :param x: 表示输入图片的特征     eg.[2, 3, 768, 768]
        :param res2: 表示xception的第一层特征输出     eg.[2, 256, 192, 192]
        :param out: 表示经过级联空洞卷积提取的特征的粗糙预测    eg.[2, 19, 48, 48]
        :return: rend:更准确的预测,points:不确定像素点的位置
        """
        """
        1. Fine-grained features are interpolated from res2 for DeeplabV3
        2. During training we sample as many points as there are on a stride 16 feature map of the input
        3. To measure prediction uncertainty
           we use the same strategy during training and inference: the difference between the most
           confident and second most confident class probabilities.
        """
        if not self.training:
            return self.inference(x, res2, out)
		#获得不确定点的坐标
        points = sampling_points(out, x.shape[-1] // 16, self.k, self.beta) #out:[2, 19, 48, 48] || x:[2, 3, 768, 768] || points:[2, 48, 2]
		#根据不确定点的坐标,得到对应的粗糙预测
        coarse = point_sample(out, points, align_corners=False) #[2, 19, 48]
        #根据不确定点的坐标,得到对应的特征向量
        fine = point_sample(res2, points, align_corners=False)  #[2, 256, 48]
		#将粗糙预测与对应的特征向量合并
        feature_representation = torch.cat([coarse, fine], dim=1)   #[2, 275, 48]
		#使用MLP进行细分预测
        rend = self.mlp(feature_representation) #[2, 19, 48]
        return {"rend": rend, "points": points}
##推理阶段
@torch.no_grad()
    def inference(self, x, res2, out):
        """
        输入:
        x:[1, 3, 768, 768],表示输入图片的特征
        res2:[1, 256, 192, 192],表示xception的第一层特征输出
        out:[1, 19, 48, 48],表示经过级联空洞卷积提取的特征的粗糙预测
        输出:
        out:[1,19,768,768],表示最终图片的预测
        主要思路:
        通过 out计算出top N = 8096 个不稳定的像素点,针对每个不稳定像素点得到在res2(fine)
        和out(coarse)中对应的特征,组合8096个不稳定像素点对应的fine和coarse得到rend,
        再通过mlp得到更准确的预测,迭代至rend的尺寸大小等于输入图片的尺寸大小
        """
        """
        During inference, subdivision uses N=8096
        (i.e., the number of points in the stride 16 map of a 1024×2048 image)
        """
        num_points = 8096
                while out.shape[-1] != x.shape[-1]: #out:[1, 19, 48, 48], x:[1, 3, 768, 768]
        	#每一次预测均会扩大2倍像素,直至与原图像素大小一致
            out = F.interpolate(out, scale_factor=2, mode="bilinear", align_corners=True)   #out[1, 19, 48, 48]
            points_idx, points = sampling_points(out, num_points, training=self.training)   #points_idx:8096 || points:[1, 8096, 2]
            coarse = point_sample(out, points, align_corners=False) #coarse:[1, 19, 8096]   表示8096个不稳定像素点根据高级特征得出的对应的类别
            fine = point_sample(res2, points, align_corners=False)  #fine:[1, 256, 8096]    表示8096个不稳定像素点根据低级特征得出的对应类别
            feature_representation = torch.cat([coarse, fine], dim=1)   #[1, 275, 8096] 表示8096个不稳定像素点合并fine和coarse的特征
            rend = self.mlp(feature_representation) #[1, 19, 8096]
            B, C, H, W = out.shape  #first:[1, 19, 128, 256]
            points_idx = points_idx.unsqueeze(1).expand(-1, C, -1)  #[1, 19, 8096]
            out = (out.reshape(B, C, -1)

python 类型注释 # type

Type Comments[类型注解]

注释是在Python 3中引入的,并且它们没有被反向移植到Python 2.这意味着如果您正在编写需要支持旧版Python的代码,则无法使用注释。

要向函数添加类型注释,您可以执行以下操作:

import math 
def circumference(radius):    
# type: (float) -> float    
   return 2 * math.pi * radius

类型注释只是注释,所以它们可以用在任何版本的Python中。

类型注释由类型检查器直接处理,所以不存在__annotations__字典对象中:

>>> circumference.__annotations__{}

类型注释必须以type: 字面量开头,并与函数定义位于同一行或下一行。如果您想用几个参数来注释一个函数,您可以用逗号分隔每个类型:

def headline(text, width=80, fill_char="-"):  
  # type: (str, int, str) -> str    
   return f" {text.title()} ".center(width, fill_char) 

print(headline("type comments work", width=40))

您还可以使用自己的注释在单独的行上编写每个参数:

# headlines.py
 
  def headline(
      text,           # type: str
      width=80,       # type: int
      fill_char="-",  # type: str
  ):                  # type: (...) -> str
      return f" {text.title()} ".center(width, fill_char)
 
 print(headline("type comments work", width=40))

通过Python和Mypy运行示例:

$  python headlines.py
---------- Type Comments Work ---------- 
$ mypy headline.py
$

如果传入一个字符串width=”full”,再次运行mypy会出现一下错误。

$ mypy headline.py
headline.py:10: error: Argument "width" to "headline" has incompatible
                       type "str"; expected "int"

您还可以向变量添加类型注释。这与您向参数添加类型注释的方式类似:

pi = 3.142  # type: float

上面的例子可以检测出pi是float类型。

A ConvNet for the 2020s

paper:https://arxiv.org/abs/2201.03545 CVPR 2022.Facebook AI Research

github:https://github.com/facebookresearch/ConvNeXt

作者提出了ConvNeXt,一个完全由标准 ConvNet 模块构建的纯 ConvNet 模型。ConvNeXt 准确、高效、可扩展且设计非常简单。

2020年以来,ViT一直是研究热点。ViT在图片分类上的性能超过卷积网络的性能,后续发展而来的各种变体将ViT发扬光大(如Swin-T,CSwin-T等),值得一提的是Swin-T中的滑窗操作类似于卷积操作,降低了运算复杂度,使得ViT可以被用做其他视觉任务的骨干网络,ViT变得更火了。本文探究卷积网络到底输在了哪里,卷积网络的极限在哪里。在本文中,作者逐渐向ResNet中增加结构(或使用trick)来提升卷积模型性能,最终将ImageNet top-1刷到了87.8%。作者认为本文所提出的网络结构是新一代(2020年代)的卷积网络(ConvNeXt),因此将文章命名为“2020年代的卷积网络”。

作者的出发点时Resnet-50模型。首先使用用于训练视觉变压器的类似训练技术训练它,与原始Resnet-50相比,获得了较大的改进效果。这将是我们的基线。然后,我们研究了一系列设计决策,总结为1)宏观设计,2)Resnext,3)inverted bottleneck,4)使用大的核,以及5)各种层的微型设计。

1、训练技巧:

作者认为,除了网络结构,训练技巧也会影响最终的效果,除了vision Transformers的结构 ,vision Transformers的一些训练技巧也给作者带来 一些启发。作者使用swin transformer的训练技巧,应用在基线模型中:

1、增加 epoch到300

2、使用AdamW优化器

3、数据增强技术(cutmix、mixup等等)

结果分类准确率由76.1%上升到78.8%。具体训练config如下:

2、宏观设计

作者借鉴了Swin-T的两个设计:

  1. 每阶段的计算量(调整每个阶段block数量)
  2. 对输入图片下采样方法

对于第一点类似Swin-T四个阶段1:1:9:1的计算量,作者将ResNet-50每个阶段block数调整为3,3,9,3(原来为3,4,6,3),增加第三阶段计算量,准确率由78.8%提升至79.4%。

这个每阶段计算量的设计:感觉很多模型都是在中间部分的计算量最多 ,两头的计算量最小,这种设计的效果最好。

对于第二点Swin-T融合压缩2×2的区域,作者则使用4×4步长为4的卷积对输入图片进行下采样,这样每次卷积操作的感受野不重叠,准确率由79.4%提升至79.5%。

3、类ResNeXt设计

depthwise conv中的逐channel卷积操作和self-attention中的加权求和很类似,因此作者采用depthwise conv替换普通卷积。参照ResNeXt,作者将通道数增加到96,准确率提升至80.5%,FLOPs相应增大到了5.3G。相比之下原始的ResNet-50 FLOPs为4G,运算量增大很多。

4、Inverted Bottleneck

在depthwise conv的基础上借鉴MobileNet的inverted bottleneck设计(维度先扩增4倍在缩减),将block由下图(a)变为(b)。因为depthwise不会使channel之间的信息交互,因此一般depthwise conv之后都会接1 × 1 × C的pointwise conv。这一顿操作下来准确率只涨了0.1%到80.6%。在后文说明的大模型上涨点多一点。

Moving up depthwise conv layer

首先,考虑到卷积核太大会导致计算复杂度上升,不方便作者去寻找大卷积核。因此作者借鉴transformer里面MSA block 放在了1*1卷积之前,把7*7的 depthwise conv layer 放在1*1卷积之前,这样, depthwise conv layer 的通道数下降,相应计算量也下降,但性能下降到了79.9%

增大卷积kernel

作者认为更大的感受野是ViT性能更好的可能原因之一,作者尝试增大卷积的kernel,使模型获得更大的感受野。首先在pointwise conv的使用上,作者为了获得更大的感受野,将depthwise conv提前到1 × 1 conv之前,之后用384个1 × 1 × 96的conv将模型宽度提升4倍,在用96个1 × 1 × 96的conv恢复模型宽度。反映在上图中就是由(b)变为(c)。由于3×3的conv数量减少,模型FLOPs由5.3G减少到4G,相应地性能暂时下降到79.9%。

然后作者尝试增大depthwise conv的卷积核大小,证明7×7大小的卷积核效果达到最佳

其他乱七八糟的尝试

借鉴最初的Transformer设计,作者将ReLU替换为GELU;ViT的K/Q/V计算中都没有用到激活函数和归一化层,于是作者也删除了大量的激活函数和归一化层,仅在1 × 1卷积之间使用激活函数,仅在7 × 7卷积和1 × 1 卷积之间使用归一化层,同时将BN升级为LN。最终block结构确定如下:

顺便复习一下各种归一化方法:

最后仿照Swin-T,作者将下采样层单独分离出来,单独使用2 × 2卷积层进行下采样。为保证收敛,在下采样后加上Layer Norm归一化。最终加强版ResNet-50准确率82.0%(FLOPs 4.5G)。

总的来说ResNet-50、本文模型和Swin-T结构差别如下:

实验结果

作者在ResNet-50加强版的基础上又提出了多个变体(ConvNeXt-T/B/L/XL),从实验结果上看ResNet-50加强版性能收益较为突出,越是大模型性能收益越低。

在检测、分割下游任务中ConvNeXt也获得了与Swin-T相似或更好的结果。结果就不细说了。

消融实验

每一部分具体涨点效果如下:

总的来说本文实验做的比较充分,总结一下,卷积网络涨点可以尝试:

  1. 对输入下采样时尝试无重叠小一点的卷积层,例如4 × 4,stride=4的卷积;
  2. block中采用大卷积核,例如7 × 7;
  3. depthwise conv + inverted bottleneck + moving up depthwise layer的block结构;
  4. 减少激活/归一化层,ReLU换成GELU,BN换成LN;
  5. 使用2×2 conv + LN下采样。

Vision MLP –Pay Attention to MLPs

MLP-Mixer的增强版,带gating的MLP。有两个版本,分别是gMLP和aMLP。Pay-Attention-to-MLPs是gMLP版本,同时也提出了gMLP的增强版aMLP。

paper: https://arxiv.org/abs/2105.08050

github: https://github.com/antonyvigouret/Pay-Attention-to-MLPs

此文和最近刊出MLP文章相同,旨在探究self-attention对于Transformer来说是否至关重要。并在CV和NLP上的相关任务进行实验。

Transformer结构具有可并行化汇聚所有token间的空间信息的优点。众所周知self-attention是通过计算输入间的空间关系动态的引入归纳偏置,同时被静态参数化的MLP能表达任意的函数,所以self-attention对于Transformer在CV和NLP等领域的成功是否是至关重要的呢?

  • 此文提出了一个基于MLP的没有self-attention结构名为gMLP,仅仅存在静态参数化的通道映射(channel projections)和空间映射(spatial projections)。同时作者通过实验发现当对空间映射的线性结果进行门机制乘法得到的效果最好
  • 此文使用gMLP做图片分类并在ImageNet上取得了与DeiT、ViT等Transformer模型相当的效果。与先前的MLP模型MLP-Mixer相比,gMLP做到了参数更少(参数减少66%)效果更强(效果提升3%)。
  • 此文使用gMLP做masked language modeling,gMLP采用和Bert一样的设置最小化perplexity取得了和Transformer模型预训练一样好的效果。通过pretraining和finetuning实验发现随着模型容量的增加,gMLP比Transformer提升更大,表明模型相较于self-attention可能对于模型容量的大小更为敏感。
  • 对于需要跨句对齐的微调任务MNLI,gMLP与Transformer相比逊色一筹。对此作者发现加上一个128特征大小的单头注意力足以使得gMLP在任何NLP任务上取得比Transformer更好的效果。

gMLP由L个如下图所示的模块堆叠而成

设每个模块的输入 \(X \in \mathbb{R}^{n \times d}\), n为序列长度, d为特征维度。每个模块表达如下:
\(Z=\sigma(X U), \quad \tilde{Z}=s(Z), \quad Y=\tilde{Z} V\)
\(\sigma\) 是GELU等激活函数, U 和 V 和Transformer中的FFN类似都是线性映射。为了简洁表达上式中 省略了shortcuts, normalizations 和 biases。
上式中最重要的是能捕捉空间交互的 \(s(\cdot)\) 。如果上式去掉 \(s(\cdot)\) 那么将不再能进行空间交互和FFN 并无区别。文中作者选择名为 Spatial Gating Unit (SGU) 的模块作为 \(s(\cdot)\) 捕捉空间依赖。另外,gMLP在NLP、CV任务中遵循与BERT、ViT一样的输入输出规则。

Spatial Gating Unit:

为了能有跨token的交互, \(s(\cdot)\) 操作须在空间维度。可以简单的使用线性映射表示:
\(f_{W, b}(Z)=W Z+b\)
其中 \(W \in \mathbb{R}^{n \times n}\) 表示空间交互的映射参数。在self-attention中 W 是通过 Z 动态计算得到的。 此文对上式使用gating操作以便更好的训练,如下所示:
\(s(Z)=Z \odot f_{W, b}(Z)\)
为了训练更稳定,作者将 W 和 b 分别初始化为接近 0 与 1 来保证在开始训练时 \(f_{W, b} \approx 1\) 、 \(s(Z z) \approx Z\) 使得在开始阶段gMLP近似于FFN并在训练中逐渐学习到跨token的空间信息。
作者进一步发现将 Z 从通道维度分割成两部分 \(\left(Z_1, Z_2\right)\) 进行gating操作更有用,如下所示:
\(
s(Z)=Z_1 \odot f_{W, b}\left(Z_2\right)
\)
另外函数 \(f_{W, b}\)的输入通常需要normalizel以此提升模型的稳定性。

一些思考:这里的SpatialGatingUnit里面用到了一个通道split,然后再将分割后的两部分做乘法,让我想到了NAFnet中的simplegate,这个的作用一是减少计算量(相比于GELU)、另外引入门控机制,在通道维度进行通道交织,对于模型的效果表现很好。

作者进一步分析了SGU与现有的一些操作的相似之处:首先是Gated Linear Units (GLU) 与 SGU的区别在于SGU对spatial dimension而GLU对channel dimension; 其次SGU和
Squeeze-and-Excite (SE) 一样使用hadamard-product,只是SGU并没有跨通道的映射来保 证排列不变性;SGU的空间映射可以看作depthwise convolution不过SGU只学习跨通道只是, 并没有跨通道过滤器;SGU学习的是二阶空间交互 \(z_i z_j\) , self-attention学习的是三阶交互 \(q_i k_j v_k\) , SGU的复杂度为 \(n^2 e / 2\) 而self-attention的复杂度为 \(2 n^2 d_{\text {。 }}\)

实验:

1、Image Classification

此文首先将gMLP应用于图片分类,使用ImageNet数据集而且不使用额外数据。下表首先展示了gMLP用于图片分类的参数,gMLP和ViT/B16一样使用 16×16 个patch,同时采用和DeiT相似的正则化方法防止过拟合。

下表中gMLP与baselines在ImageNet上的结果表示gMLP取得了与视觉Transformer相当的结果,同时与其它MLP视觉模型相比,gMLP取得了准确率、速度权衡下最好的结果。

Masked Language Modeling with BERT:

此文同时将gMLP应用于masked language modeling(MLM)任务,对于预训练和微调任务,模型的输入输出规则都保持与BERT一致。

作者观察到在MLM任务最后学习到的空间映射矩阵总是Toeplitz-like matrics,如下图所示。所以作者认为gMLP是能从数据中学习到平移不变性的概念的,这使得gMLP实质起到了卷积核是整个序列长度的1-d卷积的作用。在接下来的MLM实验中,作者初始 W 为Toeplitz matrix。

Ablation: The Importance of Gating in gMLP for BERT’s Pretraining:下表展示了gMLP的各种变体与Transoformer模型、MLP-Mixer的比较,可以看到gMLP在与Transformer相同模型大小的情况下能达到与Transformer相当的效果。同时gating操作对于空间映射十分有用。同时下图还可视化了模型学习到的空间映射参数。

Case Study: The Behavior of gMLP as Model Size Increases:下表与下图展示了gMLP随着模型增大逐渐能有与Transformer相当的效果,可见Transformer的效果应该主要是依赖于模型尺寸而非self-attention。

  • Ablation: The Usefulness of Tiny Attention in BERT’s Finetuning:从上面的Case Study可以发现gMLP对于需要跨句子连接的finetuing任务可能不及Transformer,所以作者提出了gMLP的增强版aMLP。aMLP相较于gMLP仅增加了一个单头64的self-attention如下图所示:

从下图结果可以发现aMLP相较于gMLP极大提升了效果并在所有task超过了Transformer。

Vision MLP –ResMLP

Feedforward networks for image classification with data-efficient training

我们提出了ResMLP,一种完全基于多层感知机(MLP)进行图像分类的体系结构。 它是一个简单的残差网络,它交替(i)线性层,其中图像 patches在通道之间独立且相同地交互;以及(ii)两层前馈网络,其中通道中的每个 patch独立地相互作用。

CODE:

import torch
import numpy as np
from resmlp import ResMLP

img = torch.ones([1, 3, 224, 224])

model = ResMLP(in_channels=3, image_size=224, patch_size=16, num_classes=1000,
                 dim=384, depth=12, mlp_dim=384*4)

parameters = filter(lambda p: p.requires_grad, model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('Trainable Parameters: %.3fM' % parameters)

out_img = model(img)

print("Shape of out :", out_img.shape)  # [B, in_channels, image_size, image_size]

本文作者提出了一种基于全连接层的图像分类网络。网络结构与MLP-Mixer相似,即先将输入图像拆分成若干patch,对每个patch通过全连接层转换为特征嵌入矩阵,该矩阵的两个维度分别表示channel维度(每个局部位置的特征维度)和patch维度(表示局部位置的维度)。首先将该矩阵转置后沿patch维度进行全连接层运算,实现不同patch之间的交互;再沿channel维度进行全连接运算,实现不同channel之间的交互。最后使用池化层和输出层获得分类结果。本文与MLP-Mixer的不同之处在于采用了更强的数据增强方法和蒸馏策略。

当采用现代的训练策略进行训练时,使用大量的数据增广和可选的蒸馏方法,可以在ImageNet上获得令人惊讶的良好精度/复杂度折衷。

Affine仿射变换:

函数名称:diag(x)
函数功能:构建一个n维的方阵,它的主对角线元素值取自向量x,其余元素都为0

Vision MLP系列–RepMLP

RepMLPNet: Hierarchical Vision MLP with Re-parameterized Locality (https://arxiv.org/abs/2112.11081)

CVPR 2022

RepMLP Block

Github source: https://github.com/DingXiaoH/RepMLP

最近公开了一系列视觉MLP论文,包括RepMLP、MLP-Mixer、ResMLP、gMLP等。在这个时间点出现关于MLP的一系列讨论是很合理的:

1) Transformer大火,很多研究者在拆解Transformer的过程中多多少少地对self-attention的必要性产生了疑问。去掉了self-attention,自然就剩MLP了。

2) 科学总是螺旋式上升的,“复兴”老方法(比如说另一篇“复兴”VGG的工作,RepVGG)总是喜闻乐见的。

这些论文引发了热烈的讨论,比如:

1) 这些模型到底是不是MLP?

2) 卷积和全连接(FC)的区别和联系是什么?FC是不是卷积,卷积是不是FC?

3) 真正的纯MLP为什么不行?

4) 所以MLP is all you need?

《RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition》。这篇文章讲了一个全连接层找到一份陌生的工作(直接进行feature map的变换),为了与那些已经为这份工作所特化的同胞(卷积层)们竞争,开始“内卷”的故事。

关键贡献在于,RepMLP用卷积去增强FC,既利用其全局性又赋予其局部性,并通过结构重参数化,将卷积融合到FC中去,从而在推理时去除卷积。

论文:RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition

代码:DingXiaoH/RepMLP

1. 为什么真正的纯MLP不太行?

我们一般认为多层感知机(MLP)是至少两层全连接层(FC)堆叠得到的模型,而且一般把同时含有卷积和MLP的模型(或模型中的一个模块)称为CNN。尽管目前大家对什么叫MLP的问题尚有争议(下图),我们不妨先定义一个任何人都会称之为MLP的100%纯MLP:

这个MLP在ImageNet上的输入是(3, 112, 112),第一层将其变为(32, 56, 56),第二层将其变为(64, 28, 28),然后global average pool,然后经过FC映射为1000类。这样总共只有三个FC,毫无疑问是MLP。这三层的参数为:

第一层:3x112x112x32x56x56 = 3.77G 参数

第二层:32x56x56x64x28x28 = 5.03G 参数

第三层:64×1000 = 64k 参数,忽略不计

看起来有点吓人,但这确实是一个处于A1位置的纯MLP应有的体量,虽然它只有两层,而且通道数只有32和64。除了减小通道数量,任何试图减小参数量的改动都将使其不再属于A1位置。比如说:

1) 先切块。把112×112的输入切成56×56的四块,每一块经过第一层变成28×28,再拼起来,这样第一层的参数量变成了3x56x56x32x28x28=236M,看起来好多了。但是,这破坏了全局性,因为分属于两块中的两点之间不再有联系了!换句话说,我们引入了一种局部性:一张图切成四块之后,每块中的任一像素只跟同块中的其他像素有联系。ViT,RepMLP和其他几篇MLP都用了这种操作或某种类似的操作。

2) 分组FC。正如卷积有分组卷积一样,FC也可以分组。由于torch里没有现成的算子,分组FC可以用分组1×1卷积实现。组数为g,参数量和计算量就会变成1/g。可惜,这也引入了局部性。RepMLP用了这种操作。

3) 把一个FC拆分成两次操作,第一次操作对channel维度线性重组,spatial共享参数(等价于1×1卷积);第二次操作对spatial维度线性重组,channel共享参数(等价于先转置后1×1卷积)。这思想可以类比于depthwise conv + 1×1 conv。MLP-Mixer使用这种操作,用两个各自都不具有全局性的操作实现了整体的全局性(而RepMLP使用另一种不同的机制,对不同的分块做pooling再连接,实现了这种全局性)。

所以,真正的100%纯MLP不太行,大家都在用各种花式操作做“伪MLP”的原因之一,就是体量太大。

这篇文章介绍的RepMLP属于B2的位置,不追求纯MLP。称其为“MLP”的原因是想强调卷积和FC的区别:RepMLP将卷积看成一种特殊的FC,显式地用卷积去强化FC(把FC变得具有局部性又不失全局性),指出了这样的FC强在哪里(如ResNet-50中,用一半通道数量的RepMLP替换3×3卷积就可以实现同等精度和55%加速),并用这种强化过的FC(及一些其他技巧)构造一种通用的CNN基本组件,提升多重任务性能。论文中说明了这里MLP的意思是推理时结构“不包含大于1×1的卷积”。

2. RepMLP:FC“内卷”,卷出性能

真正的100%纯MLP不太行的原因之二,是不具有局部先验。

在一张图片中,一个像素点跟它周围的像素点的关系往往比远在天边的另一个像素点更密切,这称为局部性。人类在识别图片的时候潜意识地利用这一点,称为局部先验。卷积网络符合局部先验,因为卷积核通过滑动窗口在图片上“一块一块地”寻找某种特征。

那么FC层呢?FC能自动学到这一点吗?在有限的数据量(ImageNet)和有限的计算资源前提(GPU)下,很难。

实验验证:下面我们假设FC层的输入是64x10x10的feature map直接 “展平”成的6400维向量。输出也是6400维向量,然后reshape成64x10x10的feature map。下图展示了FC学得的kernel中的一个切片的权值大小。简单地讲(详见论文),展示的这一部分表示在输出的第0个channel中随便找的一个采样点(6,6)(也就是图中黄框标出来的点)作用于第0个输入channel上的10×10个像素点的权值。颜色越深,表示权值越大。比如说,如果图中的(5,5)点颜色深,就表示这个FC层认为输出中的(0,6,6)点与输入中的(0,5,5)点关系紧密。

结果很明显,(6,6)周围的权值并没有颜色更深,也就是说FC并不认为这个点和周围点的联系更紧密。相反,似乎这个FC层认为(6,6)点与右上和右下部分关系更密切。实验也证明,不具有局部性的FC效果较差。

既然图像的局部性很强,FC把握不住,那怎么办呢?RepMLP提出,用卷积去增强FC(如下图所示,输入既被展平成向量并输入FC,又用不同大小的卷积核进行卷积,各自过BN后相加),并通过结构重参数化,将卷积融合到FC中去,从而在推理时去除卷积。

我们将卷积和FC之间建立联系,是因为卷积可以看成一个稀疏且存在重复参数的FC。如下图代码所示,给定输入X和卷积核conv_K,其卷积的结果等于X(直接展平成向量)和fc_K的矩阵乘,fc_K称为conv_K的等效FC核。尽管我们都相信这样的fc_K一定存在,但根据conv_K的值直接构造出fc_K的方法(下图中的convert_K函数)似乎不太简单。

本文提出了一种简洁优美的做法(见后文)。我们用这种方法构造出fc_K并打印出来,可以看出它是一个稀疏且有很多元素相同的矩阵(Toeplitz矩阵)。如下图的代码和结果所示。

RepMLP把卷积的输出和FC的输出相加,这样做的好处是:

1) 降低FLOPs,提高速度。用我们提出的方法把卷积全都转换为等效FC kernel后,由于矩阵乘法的可加性(AX + BX = (A+B)X),一个稀疏且共享参数的FC(Toeplitz矩阵)加一个不稀疏不共享参数的FC(全自由度的矩阵),可以等价转换为一个FC(其参数是这两个矩阵之和)。这样我们就可以将这些卷积等效地去掉。这一思路也属于结构重参数化(通过参数的等价转换实现结构的等价转换,如RepVGG)。

2) 在同等参数量的情况下,FC的FLOPs远低于卷积。

3) 相比于纯FC,这样做产生了局部性。注意这种局部性是我们“赋予”FC的,而不是让FC学到的。

4) 相比于卷积层,这样做使得相距遥远的两个点直接相连,具备了全局性。

这样做看起来像是让FC的“内部”含有卷积,所以也可以称为“内卷”。事实证明,跟人类相似,FC的“内卷”也可以提高性能。

只剩下一个问题了:我们相信存在一个FC kernel等价于卷积的卷积核,但是给定一个训练好的卷积核,怎么构造出FC kernel(Toeplitz矩阵)呢?

其实也很简单:FC kernel等于在单位矩阵reshape成的feature map上用卷积核做卷积的结果。这一做法是高效、可微、与具体的卷积算法和平台无关的。推导过程也很简洁(详见论文)。

现在,整个流程就很清晰了:

1) 训练时,既有FC又有卷积,输出相加。

2) 训练完成后,先把BN的参数“吸”到卷积核或FC中去(跟RepVGG一样),然后把每一个卷积转换成FC,把所有FC加到一起。从此以后,不再有卷积,只有FC。

3) 保存并部署转换后的模型。

现在我们再看一下用卷积增强后转换得到的FC kernel,可以看出采样点周围的权值变大了,现在(6,6)点更关注它旁边的输入点了。有趣的是,这里用到的最大卷积是7×7,但是7×7的范围(蓝色框)外还有一些值(红色框)比蓝框内的值大,这说明全局性也没有被局部性“淹没”。

一些其他设计

RepMLP中也用了一些其他设计,包括:

1) 用groupwise conv实现groupwise FC,减少参数和计算量。

2) 将输入分块(最近大家都会用的常见操作),进一步减少参数和计算量。如下图所示,H和W是feature map的分辨率,h和w是每一块的分辨率。

3) 用两个FC在不同分块之间建立联系,确保全局性。如下图所示。

实验结果

用RepMLP替换Res50中的部分结构,在ImageNet上有性能提升。将ImageNet pretrained模型迁移到语义分割和人脸上,也都有性能提升。

在ImageNet上的实验是在Res50中做的。考虑到Res50的主干通道较多(256、512、1024、2048),为了将RepMLP用到Res50中取得合理的trade-off,我们做了以下设计:

1)RepMLP Bottleneck Block:在RepMLP之前用1×1和3×3降维,RepMLP之后用3×3和1×1升维。这一结构类似于旷视在工程中探索并申请的专利GLFP(202010422194.X, Visual task processing method and device and electronic system,下图)。

2)RepMLP Light Block:在RepMLP之前用1×1大幅降维,之后用1×1大幅升维。降维/升维的幅度(8x)比Res50(4x)更大。

一些有趣的发现:

1) RepMLP中具有局部先验的成分(融合进FC的卷积),所以对于具有平移不变性的任务(ImageNet,Cityscapes语义分割)有效。

2) RepMLP中也具有不具有平移不变性的成分(大FC kernel),所以对于具有某种位置模式(例如人脸图像中,眼睛总是在鼻子上面)的任务也有效。

3) 由于FC和卷积的差别,RepMLP可以大幅增加参数而不降低速度(参数增加47%,ImageNet精度提升0.31%,速度仅降低2.2%)。

一些常见问题

RepMLP和ResMLP是什么关系?

相当于旺旺碎冰冰和王冰冰的关系。只是名字有点像。RepMLP中用卷积增强FC的思路也可以用在其他MLP架构中,应该也会有提升。另外,ResMLP、RepMLP和ResRep(去年做的一篇用重参数化做剪枝的论文)也没有关系。

把卷积融合进FC里,那FC不就是卷积了吗?

卷了,但不是完全卷,而且比卷积更强。上面可视化的图显示,转换后的kernel可以关注到卷积核的感受野以外的信息,因而表征能力更强。论文中报告的实验表明,这样的操作可以以一半的channel量达到与纯CNN相当的性能,速度更快,FLOPs更低。本文的关键也在于把卷积看成一种特殊的FC,然后考虑如何利用这种特殊性

所以MLP is all you need?

目前看来,还差得远。目前的方法多多少少都用到了切块等操作,都需要用某种方式降低参数量和引入局部性。真正的纯MLP(A1位置)依然还没有希望。真正纯MLP的一个大麻烦是总的参数量和输入分辨率耦合,因而改变输入分辨率会很困难。MLP-Mixer的一个缺点是不方便改变输入分辨率,所以它在ImageNet分类上的性能不容易迁移到其他任务上去。

ps–进程查看器

写这个的原因:服务器多人共享使用,有些时候,有些进程可能不知道是哪位大哥用户,而某些进程可能对于服务器来说需要kill掉,这时候就需要查看进程的详细信息(比如,跑 GPU,有些用户虽然程序停止了,但去后台发现其实还在占用显存……,这时候就需要去查看当前进程的用户 是哪个,然后kill)

查看PID对应的用户的方法:ps aux | grep [your_PID]

Linux中的ps命令是Process Status的缩写。ps命令用来列出系统中当前运行的那些进程。ps命令列出的是当前那些进程的快照,就是执行ps命令的那个时刻的那些进程,如果想要动态的显示进程信息,就可以使用top命令/htop命令。

要对进程进行监测和控制,首先必须要了解当前进程的情况,也就是需要查看当前进程,而 ps 命令就是最基本同时也是非常强大的进程查看命令。使用该命令可以确定有哪些进程正在运行和运行的状态、进程是否结束、进程有没有僵死、哪些进程占用了过多的资源等等。总之大部分信息都是可以通过执行该命令得到的。

ps 为我们提供了进程的一次性的查看,它所提供的查看结果并不动态连续的;如果想对进程时间监控,应该用 top linux下的任务管理器 工具。

注:kill 命令用于杀死进程。

linux上进程有5种状态:

  1. 运行(正在运行或在运行队列中等待)
  2. 中断(休眠中, 受阻, 在等待某个条件的形成或接受到信号)
  3. 不可中断(收到信号不唤醒和不可运行, 进程必须等待直到有中断发生)
  4. 僵死(进程已终止, 但进程描述符存在, 直到父进程调用wait4()系统调用后释放)
  5. 停止(进程收到SIGSTOP, SIGTSTP, SIGTTIN, SIGTTOU信号后停止运行运行)

ps工具标识进程的5种状态码:

  • D 不可中断 uninterruptible sleep (usually IO)
  • R 运行 runnable (on run queue)
  • S 中断 sleeping
  • T 停止 traced or stopped
  • Z 僵死 a defunct (”zombie”) process

命令参数:

  • a 显示所有进程
  • -a 显示同一终端下的所有程序
  • -A 显示所有进程
  • c 显示进程的真实名称
  • -N 反向选择
  • -e 等于“-A”
  • e 显示环境变量
  • f 显示程序间的关系
  • -H 显示树状结构
  • r 显示当前终端的进程
  • T 显示当前终端的所有程序
  • u 指定用户的所有进程
  • -au 显示较详细的资讯
  • -aux 显示所有包含其他使用者的行程
  • -C<命令> 列出指定命令的状况
  • –lines<行数> 每页显示的行数
  • –width<字符数> 每页显示的字符数
  • –help 显示帮助信息
  • –version 显示版本显示

执行ps -aux 的输出:-aux 显示所有包含其他使用者的行程

  • F 代表这个程序的旗标 (flag), 4 代表使用者为 super user
  • S 代表这个程序的状态 (STAT),关于各 STAT 的意义将在内文介绍
  • UID 程序被该 UID 所拥有
  • PID 进程的ID
  • PPID 则是其上级父程序的ID
  • C CPU 使用的资源百分比
  • PRI 这个是 Priority (优先执行序) 的缩写,详细后面介绍
  • NI 这个是 Nice 值,在下一小节我们会持续介绍
  • ADDR 这个是 kernel function,指出该程序在内存的那个部分。如果是个 running的程序,一般就是 “-“
  • SZ 使用掉的内存大小
  • WCHAN 目前这个程序是否正在运作当中,若为 – 表示正在运作
  • TTY 登入者的终端机位置
  • TIME 使用掉的 CPU 时间。
  • CMD 所下达的指令为何

 ps 与grep 组合使用,查找特定进程

显示指定用户信息:

PyTorch医学图像分割开源库

github: https://github.com/MontaEllis/Pytorch-Medical-Segmentation

基于PyTorch的专注于医学图像分割的开源库,其支持模型丰富,方便易用。其可算为torchio的一个实例,作者将其综合起来,包含众多经典算法,实用性比较强。

该库特点:

  1. 支持2D和3D医学图像分割,可以修改hparam.py文件来确定是2D分割还是3D分割以及是否可以进行多分类。
  2. 支持绝大数主流分割模型,几乎提供了所有的2D和3D分割的算法。
  3. 兼容几乎所有的医学数据格式(例如 nii.gz, nii, mhd, nrrd, …),修改hparam.py的fold\_arch即可。

作者提供了训练和测试推断的代码,简单配置后训练和推断都仅需要一行命令。

已包含的分割模型:

AI部署系列:你知道模型权重的小秘密吗???

今天简单聊聊模型权重,也就是我们俗称的weight

深度学习中,我们一直在训练模型,通过反向传播求导更新模型的权重,最终得到一个泛化能力比较强的模型。同样,如果我们不训练,仅仅随机初始化权重,同样能够得到一个同样大小的模型。虽然两者大小一样,不过两者其中的权重信息分布相差会很大,一个脑子装满了知识、一个脑子都是水,差不多就这个意思。

所谓的AI模型部署阶段,说白了就是将训练好的权重挪到另一个地方去跑。一般来说,权重信息以及权重分布基本不会变(可能会改变精度、也可能会合并一些权重)。

不过执行模型操作(卷积、全连接、反卷积)的算子会变化,可能从Pytorch->TensorRT或者TensorFlow->TFLITE,也就是实现算子的方式变了,同一个卷积操作,在Pytorch框架中是一种实现,在TensorRT又是另一种时间,两者的基本原理是一样的,但是精度和速度不一样,TensorRT可以借助Pytorch训练好的卷积的权重,实现与Pytorch中一样的操作,不过可能更快些。

权重/Weight/CheckPoint

那么权重都有哪些呢?他们长什么样?

这还真不好描述…其实就是一堆数据。对的,我们千辛万苦不断调优训练出来的权重,就是一堆数据而已。也就是这个神奇的数据,搭配各种神经网络的算子,就可以实现各种检测、分类、识别的任务。

例如上图,我们用Netron这个工具去查看某个ONNX模型的第一个卷积权重。很显然这个卷积只有一个W权重,没有偏置b。而这个卷积的权重值的维度是[64,3,7,7],也就是输入通道3、输出通道64、卷积核大小7x7

再仔细看,其实这个权重的数值范围相差还是很大,最大的也就0.1的级别。但是最小的呢,肉眼看了下(其实应该统计一波),最小的竟然有1e-10级别。

一般我们训练的时候,输入权重都是0-1,当然也有0-255的情况,但不论是0-1还是0-255,只要不溢出精度上限和下限,就没啥问题。对于FP32来说,1e-10是小case,但是对于FP16来说就不一定了。

我们知道FP16的普遍精度是~5.96e−8 (6.10e−5) … 65504,具体的精度细节先不说,但是可以很明显的看到,上述的1e-10的精度,已经溢出了FP16的精度下限。如果一个模型中的权重分布大部分都处在溢出边缘的话,那么模型转换完FP16精度的模型指标可能会大大下降。

除了FP16,当然还有很多其他精度(TF32、BF16、IN8),这里暂且不谈,不过有篇讨论各种精度的文章可以先了解下。

话说回来,我们该如何统计该层的权重信息呢?利用Pytorch中原生的代码就可以实现:

# 假设v是某一层conv的权重,我们可以简单通过以下命令查看到该权重的分布
v.max()
tensor(0.8559)
v.min()
tensor(-0.9568)
v.abs()
tensor([[0.0314, 0.0045, 0.0182,  ..., 0.0309, 0.0204, 0.0345],
        [0.0295, 0.0486, 0.0746,  ..., 0.0363, 0.0262, 0.0108],
        [0.0328, 0.0582, 0.0149,  ..., 0.0932, 0.0444, 0.0221],
        ...,
        [0.0337, 0.0518, 0.0280,  ..., 0.0174, 0.0078, 0.0010],
        [0.0022, 0.0297, 0.0167,  ..., 0.0472, 0.0006, 0.0128],
        [0.0631, 0.0144, 0.0232,  ..., 0.0072, 0.0704, 0.0479]])
v.abs().min() # 可以看到权重绝对值的最小值是1e-10级别
tensor(2.0123e-10)
v.abs().max()
tensor(0.9568)
torch.histc(v.abs()) # 这里统计权重的分布,分为100份,最小最大分别是[-0.9558,0.8559]
tensor([3.3473e+06, 3.2437e+06, 3.0395e+06, 2.7606e+06, 2.4251e+06, 2.0610e+06,
        1.6921e+06, 1.3480e+06, 1.0352e+06, 7.7072e+05, 5.5376e+05, 3.8780e+05,
        2.6351e+05, 1.7617e+05, 1.1414e+05, 7.3327e+04, 4.7053e+04, 3.0016e+04,
        1.9576e+04, 1.3106e+04, 9.1220e+03, 6.4780e+03, 4.6940e+03, 3.5140e+03,
        2.8330e+03, 2.2040e+03, 1.7220e+03, 1.4020e+03, 1.1130e+03, 1.0200e+03,
        8.2400e+02, 7.0600e+02, 5.7900e+02, 4.6400e+02, 4.1600e+02, 3.3400e+02,
        3.0700e+02, 2.4100e+02, 2.3200e+02, 1.9000e+02, 1.5600e+02, 1.1900e+02,
        1.0800e+02, 9.9000e+01, 6.9000e+01, 5.2000e+01, 4.9000e+01, 2.2000e+01,
        1.8000e+01, 2.8000e+01, 1.2000e+01, 1.3000e+01, 8.0000e+00, 3.0000e+00,
        4.0000e+00, 3.0000e+00, 1.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00,
        1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 2.0000e+00,
        0.0000e+00, 2.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00,
        2.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00])

这样看如果觉着不是很直观,那么也可以自己画图或者通过Tensorboard来时候看。

那么看权重分布有什么用呢?

肯定是有用处的,训练和部署的时候权重分布可以作为模型是否正常,精度是否保持的一个重要信息。不过这里先不展开说了。

有权重,所以重点关照

在模型训练过程中,有很多需要通过反向传播更新的权重,常见的有:

  • 卷积层
  • 全连接层
  • 批处理化层(BN层、或者各种其他LN、IN、GN)
  • transformer-encoder层
  • DCN层

这些层一般都是神经网络的核心部分,当然都是有参数的,一定会参与模型的反向传播更新,是我们在训练模型时候需要注意的重要参数。

# Pytorch中conv层的部分代码,可以看到参数的维度等信息
self._reversed_padding_repeated_twice = _reverse_repeat_tuple(self.padding, 2)
if transposed:
    self.weight = Parameter(torch.Tensor(
        in_channels, out_channels // groups, *kernel_size))
else:
    self.weight = Parameter(torch.Tensor(
        out_channels, in_channels // groups, *kernel_size))
if bias:
    self.bias = Parameter(torch.Tensor(out_channels))

也有不参与反向传播,但也会随着训练一起更新的参数。比较常见的就是BN层中的running_meanrunning_std

# 截取了Pytorch中BN层的部分代码
def __init__(
    self,
    num_features: int,
    eps: float = 1e-5,
    momentum: float = 0.1,
    affine: bool = True,
    track_running_stats: bool = True
) -> None:
    super(_NormBase, self).__init__()
    self.num_features = num_features
    self.eps = eps
    self.momentum = momentum
    self.affine = affine
    self.track_running_stats = track_running_stats
    if self.affine:
        self.weight = Parameter(torch.Tensor(num_features))
        self.bias = Parameter(torch.Tensor(num_features))
    else:
        self.register_parameter('weight', None)
        self.register_parameter('bias', None)
    if self.track_running_stats:
        # 可以看到在使用track_running_stats时,BN层会更新这三个参数
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
    else:
        self.register_parameter('running_mean', None)
        self.register_parameter('running_var', None)
        self.register_parameter('num_batches_tracked', None)
    self.reset_parameters()

可以看到上述代码的注册区别,对于BN层中的权重和偏置使用的是register_parameter,而对于running_meanrunning_var则使用register_buffer,那么这两者有什么区别呢,那就是注册为buffer的参数往往不会参与反向传播的计算,但仍然会在模型训练的时候更新,所以也需要认真对待。

关于BN层,转换模型和训练模型的时候会有暗坑,需要注意一下。

刚才描述的这些层都是有参数的,那么还有一些没有参数的层有哪些呢?当然有,我们的网络中其实有很多op,仅仅是做一些维度变换、索引取值或者上/下采样的操作,例如:

  • Reshape
  • Squeeze
  • Unsqueeze
  • Split
  • Transpose
  • Gather

等等等等,这些操作没有参数仅仅是对上一层传递过来的张量进行维度变换,用于实现一些”炫技“的操作。至于这些炫技吗,有些很有用有些就有些无聊了。

上图这一堆乱七八槽的op,如果单独拆出来都认识,但是如果都连起来(像上图这样),估计连它爸都不认识了。

开个玩笑,其实有时候在通过Pytorch转换为ONNX的时候,偶尔会发生一些转换诡异的情况。比如一个简单的reshape会四分五裂为gather+slip+concat,这种操作相当于复杂化了,不过一般来说这种情况可以使用ONNX-SIMPLIFY去优化掉,当然遇到较为复杂的就需要自行优化了。

哦对了,对于这些变形类的操作算子,其实有些是有参数的,例如下图的reshap:

像这种的op,怎么说呢,有时候会比较棘手。如果我们想要将这个ONNX模型转换为TensorRT,那么100%会遇到问题,因为TensorRT的解释器在解析ONNX的时候,不支持reshape层的shape是输入TensorRT,而是把这个shape当成attribute来处理,而ONNX的推理框架Inference则是支持的。

不过这些都是小问题,大部分情况我们可以通过改模型或者换结构解决,而且成本也不高。但是还会有一些其他复杂的问题,可能就需要我们重点研究下了。

提取权重

想要将训练好的模型从这个平台部署至另一个平台,那么首要的就是转移权重。不过实际中大部分的转换器都帮我们做好了(比如onnx-TensorRT),不用我们自己操心!

不过如果想要对模型权重的有个整体认知的话,还是建议自己亲手试一试。

Caffe2Pytorch

先简单说下Caffe和Pytorch之间的权重转换。这里推荐一个开源仓库Caffe-python,已经帮我们写好了提取Caffemodel权重和根据prototxt构建对应Pytorch模型结构的过程,不需要我们重复造轮子。

我们都知道Caffe的权重使用Caffemodel表示,而相应的结构是prototxt。如上图,左面是prototxt右面是caffemodel,而caffemodel使用的是protobuf这个数据结构表示的。我们当然也要先读出来:

model = caffe_pb2.NetParameter()
print('Loading caffemodel: ' + caffemodel)
with open(caffemodel, 'rb') as fp:
    model.ParseFromString(fp.read())

caffe_pb2就是caffemodel格式的protobuf结构,具体的可以看上方老潘提供的库,总之就是定义了一些Caffe模型的结构。

而提取到模型权重后,通过prototxt中的模型信息,挨个从caffemodel的protobuf权重中找,然后复制权重到Pytorch端,仔细看这句caffe_weight = torch.from_numpy(caffe_weight).view_as(self.models[lname].weight),其中self.models[lname]就是已经搭建好的对应Pytorch的卷积层,这里取weight之后通过self.models[lname].weight.data.copy_(caffe_weight)将caffe的权重放到Pytorch中。

很简单吧。

if ltype in ['Convolution', 'Deconvolution']:
    print('load weights %s' % lname)
    convolution_param = layer['convolution_param']
    bias = True
    if 'bias_term' in convolution_param and convolution_param['bias_term'] == 'false':
        bias = False
    # weight_blob = lmap[lname].blobs[0]
    # print('caffe weight shape', weight_blob.num, weight_blob.channels, weight_blob.height, weight_blob.width)
    caffe_weight = np.array(lmap[lname].blobs[0].data)
    caffe_weight = torch.from_numpy(caffe_weight).view_as(self.models[lname].weight)
    # print("caffe_weight", caffe_weight.view(1,-1)[0][0:10])
    self.models[lname].weight.data.copy_(caffe_weight)
    if bias and len(lmap[lname].blobs) > 1:
        self.models[lname].bias.data.copy_(torch.from_numpy(np.array(lmap[lname].blobs[1].data)))
        print("convlution %s has bias" % lname)

Pytorch2TensorRT

先举个简单的例子,一般我们使用Pytorch模型进行训练。训练得到的权重,我们一般都会使用torch.save()保存为.pth的格式。

PTH是Pytorch使用python中内置模块pickle来保存和读取,我们使用netron看一下pth长什么样。。

可以看到只有模型中有参数权重的表示,并不包含模型结构。不过我们可以通过.py的模型结构一一加载.pth的权重到我们模型中即可。

看一下我们读取.pth后,state_dictkey。这些key也就对应着我们在构建模型时候注册每一层的权重名称和权重信息(也包括维度和类型等)。

当然这个pth也可以包含其他字符段{'epoch': 190, 'state_dict': OrderedDict([('conv1.weight', tensor([[...,比如训练到多少个epoch,学习率啥的。

对于pth,我们可以通过以下代码将其提取出来,存放为TensorRT的权重格式。

def extract_weight(args):
    # Load model
    state_dict = torch.load(args.weight)
    with open(args.save_path, "w") as f:
        f.write("{}\n".format(len(state_dict.keys())))
        for k, v in state_dict.items():
            vr = v.reshape(-1).cpu().numpy()
            f.write("{} {} ".format(k, len(vr)))
            for vv in vr:
                f.write(" ")
                f.write(struct.pack(">f", float(vv)).hex())
            f.write("\n")

需要注意,这里的TensorRT权重格式指的是在build之前的权重,TensorRT仅仅是拿来去构建整个网络,将每个解析到的层的权重传递进去,然后通过TensorRT的network去build好engine

// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
    std::cout << "Loading weights: " << file << std::endl;
    std::map<std::string, Weights> weightMap;

    // Open weights file
    std::ifstream input(file);
    assert(input.is_open() && "Unable to load weight file.");

    // Read number of weight blobs
    int32_t count;
    input >> count;
    assert(count > 0 && "Invalid weight map file.");

    while (count--)
    {
        Weights wt{DataType::kFLOAT, nullptr, 0};
        uint32_t size;

        // Read name and type of blob
        std::string name;
        input >> name >> std::dec >> size;
        wt.type = DataType::kFLOAT;

        // Load blob
        uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
        for (uint32_t x = 0, y = size; x < y; ++x)
        {
            input >> std::hex >> val[x];
        }
        wt.values = val;
        wt.count = size;
        weightMap[name] = wt;
    }
    std::cout << "Finished Load weights: " << file << std::endl;
    return weightMap;
}

那么被TensorRT优化后?模型又长什么样子呢?我们的权重放哪儿了呢?

肯定在build好后的engine里头,不过这些权重因为TensorRT的优化,可能已经被合并/移除/merge了。

模型参数的学问还是很多,近期也有很多相关的研究,比如参数重参化,是相当solid的工作,在很多训练和部署场景中经常会用到。