一、介绍
-
论文链接[2401.06426] UPDP: A Unified Progressive Depth Pruner for CNN and Vision Transformer (arxiv.org)
-
UPDP 是 AMD AEAI 发表的一篇做模型压缩的论文,该方法通过模块替换、插入,渐进性训练和重参数技术,能对 CNN 和 Vision Transformer 模型都进行压缩,并取得 SOTA 的压缩效果。
-
模型优化技术通常有模型剪枝,模型量化和 NAS 等技术,其中模型剪枝是一种非常重要的模型优化技术。模型剪枝通常可以分为 channel-wise 和 layer-wise 剪枝方法。channle-wise 剪枝方法通常分析出 Conv/MLP 层中不重要的 weights,通过 mask 置零这些不重要的权重和 retrain 恢复精度,最终得到剪枝的 model,layer-wise 方法则是比较粗粒度的直接剪掉不重要的 layer 或 block,最终通过 retrain 恢复精度然后得到剪枝的 model。
-
目前剪枝算法存在一些问题:channel-wise 剪枝算法对于有 depth-wise Conv 的模型如 MobileNet 系列,ConvNext 系列很难取得很好的效果。layer-wise 剪枝算法如 DepthShrinker 和 LayerFold 只在 MobileNet 和 ResNet 模型上验证过,同时这些方法是没法压缩含非BN 的 Norm Layer model,那当然没法压缩有 LayerNorm 的 Vision Transformer models. 同时这些方法在 retrain 时,通常直接粗暴去掉对应 ReLU 层导致精度无法极大恢复。
-
AMD UPDP 则是针对这些问题提出了有效的解决方法:用 BN 替换 LN/GN,用 BN 替换复杂激活函数Swish,以及渐进式的子网训练策略和重参数策略,故而 UPDP 既能压缩基本所有 CNN model 也能压缩 Vision Transformer 模型,并在这些模型取得了 SOTA 效果。
二、方法介绍
-
UPDP overview:包含超网训练,子网搜索,子网训练和子网合并。
-
超网训练:如下图所示,基于 Baseline 模型 block,构建 supernet block;supernet block 包含了 baseline block 和 pruned block flow,其相同部分share weights,不同部分则通过 lambda 因子来控制 baseline block 和 pruned block 所占比例(两者比例之和为1)。如果 baseline block 为 1, 则为 baseline block,如果pruned block 为 1,则为 pruned block。超网训练通过 sandwich rule 选择一个 baseline model,一个 random pruned model 和一个全 pruned model 一起梯度更新,这样能保证当前非充分训练的子网精度和最终充分训练子网精度的强相关性,这种方法在 NAS 中已经普遍采用。这样超网训练的时间可以不用花费太多。
-
子网搜索:在完成 step 1 超网训练后,我们需要在指定剪枝 block 数的约束下,找到一个精度最高的子网。在第一步已经说过,在 supernet block 中,lambda 置 0 表示不剪枝,选择 baseline block;置1表示剪枝,选择 pruned block,所以就是在supernet 中找出指定数目的1。通过遗传算法比较容易找到一个较优的子网。
-
子网训练:在找出最优子网后,直接将模型切换成子网,直接 retrain,将无法极大地恢复精度,UPDP 提出渐进式训练方法,将 baseline model 平滑地过渡到子网,这样能极大地保护原来的权重,从而能恢复更多的精度。
-
子网合并:在得到最终的子网后,如 Fig2 所示,原本 block 中的 ReLU 被去掉,LN 换成 BN,在 Vision Transformer 中 GELU 换成 BN等,这样就为重参数化创建了条件,相邻 2个 Conv/ FC 合并成1个 Conv/FC,BN 合并进 Conv/FC,以及 Skip Connet 也能合入绕过的单 Conv/FC 层。
-
Note:对于 Plain Model,可以自己定义 block,如将连续的几层看成一个 Block,以合并连续的 Conv 层。
-
三、实验结果
-
SOTA 方法对比:UPDP 在 ResNet34,MoblieNetV2, ConvNextV1 和 DeiT 均取得 SOTA 效果,验证方法的有效性。
-
Table 6 消融实验证明了渐进式训练的有效性。
-
四、一些思考
-
UPDP 确实是一篇不错的压缩工作,将深度压缩做的通用且有效。
-
Vision Transformer model 压缩目前的性能不是最强的,仍值得结合蒸馏等方法进一步探索。
-
其中模型替换思想是非常值得扩展的,如用简单的 OP(BN,Linear)替换复杂的 OP,辅之以渐进式的训练方法极大的恢复精度,其中可以尝试用 ReLU 替换 Multi-Head Attention 的 Softmax Op