Universally Slimmable Networks and Improved Training Techniques 解读

一、论文基本信息论文题目Universally Slimmable Networks and Improved Training Techniques作者Jiahui Yu、Thomas Huang发表会议ICCV 2019方法简称US-NetUniversally Slimmable Network论文链接CVF Open Access / arXiv官方代码JiahuiYu/slimmable_networks这篇论文的核心目标是训练一个“宽度可任意调节的单一网络”。普通模型训练完成后宽度固定静态剪枝后也只能得到一个固定小模型而 US-Net 希望一个模型在运行时可以根据资源预算选择不同宽度。论文明确提出US-Net 将原来的 slimmable network 从“只能在预定义宽度集合中切换”扩展到arbitrary width并提出sandwich rule和inplace distillation来提升训练效果。(arXiv)这篇论文不是传统意义上的通道剪枝方法。它不会训练完后永久删除某些通道而是训练一个共享权重的超网络运行时根据宽度比例截取前若干通道来形成不同子网络。官方代码与模型也由作者开源在JiahuiYu/slimmable_networks。(arXiv)二、论文要解决的问题深度模型部署时经常会遇到一个现实问题不同设备的计算能力差别很大同一个模型可能要部署到服务器、高端手机、低端移动端、摄像头、嵌入式设备等环境中。如果每一种计算预算都单独训练一个模型训练、调参、存储和维护成本都会很高。传统压缩方法通常有两种思路。一种是静态压缩例如通道剪枝、filter pruning、低秩分解、量化等。这类方法会得到一个固定模型。问题是一个固定模型只能对应一个固定计算预算。另一种是动态推理例如根据输入样本动态跳过 block 或通道。这类方法可以根据样本难度改变计算量但通常涉及 gate、policy network、强化学习或动态稀疏卷积实现复杂度较高。US-Net 关注的是另一个问题能不能只训练一个模型让它在运行时根据资源预算自由切换宽度早期 Slimmable Neural Networks 已经证明一个网络可以在几个离散宽度之间切换。但它的问题是宽度集合是预先固定的不能任意选择中间宽度。Slimmable Neural Networks 的核心是训练一个可在多个宽度下运行的共享模型并使用 switchable batch normalization 解决不同宽度的 BN 统计差异。(arXiv)US-Net 进一步提出不应该只支持几个固定宽度而应该支持一个范围内的任意宽度。例如设备预算刚好适合那么模型就应该可以直接以宽度运行而不必重新训练一个模型。三、核心思想US-Net 的核心思想可以概括为一句话把一个宽网络看成由多个窄网络嵌套组成的共享权重模型运行时只使用前 (k) 个通道就得到一个对应宽度的子网络。假设某一层完整宽度有 (n) 个通道如果运行宽度比例为 (r)那么实际使用通道数大致为当 r1.0 时使用完整网络。当 r0.5 时只使用前一半通道。当 r0.25时只使用前四分之一通道。这种设计的关键在于所有宽度共享同一套权重。也就是说子网络、子网络和子网络不是三个独立模型而是同一个模型中不同宽度的子结构。小宽度网络使用前面的通道大宽度网络在小宽度基础上继续增加后面的通道。论文用部分通道聚合来解释这种嵌套关系。完整输出可以写成如果只使用前 (k) 个通道则输出为随着 (k) 增大会逐渐接近完整输出。因此更宽的子网络可以看作在更窄子网络的基础上继续补充表达能力。这个视角解释了为什么一个共享权重网络有可能支持连续宽度运行。这里最需要强调的是US-Net 的目标不是得到一个最小模型而是得到一个可以覆盖多个计算预算的模型族。四、US-Net 与普通剪枝的区别US-Net 很容易被误解为一种通道剪枝方法但它和传统剪枝有明显区别。传统通道剪枝通常会永久删除一部分通道得到一个固定结构的小网络。剪枝完成后所有输入都使用同一个小模型。US-Net 不永久删除通道而是保留完整模型。运行时根据宽度比例选择前多少通道参与计算。所以传统剪枝更像是训练一个模型然后剪成一个固定小模型。US-Net 更像是训练一个共享权重的大模型同时让它内部所有宽度的子网络都可用。这带来一个非常重要的优势一个 US-Net 可以提供连续的 accuracy–efficiency trade-off。部署时不必为不同设备保存多个模型只需要选择合适宽度即可。五、关键问题一任意宽度下的 BatchNormUS-Net 最大的技术难点之一是 BatchNorm。在普通网络中BN 在训练阶段使用 mini-batch 统计量在推理阶段使用训练期间累积的 running mean 和 running variance。问题是不同宽度的子网络激活分布不同因此 BN 的均值和方差也不同。如果只支持几个固定宽度可以为每个宽度保存一套 BN 统计量。这就是早期 Slimmable Network 中的Switchable Batch Normalization。但 US-Net 支持任意宽度如果每个宽度都保存一套 BN数量会非常大几乎不可行。论文给出的解决方案是BN post-statistics。核心做法是训练阶段不为每个宽度单独维护 BN 统计量训练完成后固定网络权重再为目标宽度重新统计 BN 均值和方差。具体过程是选择一个目标宽度用少量训练样本前向传播重新计算该宽度下 BN 的 mean 和 variance然后用于推理。这一步很关键因为它让 US-Net 不需要在训练时存储无数套 BN 参数。论文实验也表明重新统计 BN 并不需要完整训练集少量图像就能得到接近完整统计的效果。因此US-Net 支持任意宽度的第一个基础是用训练后 BN 重新统计替代训练时保存所有宽度的 BN。六、关键问题二如何高效训练所有宽度如果一个网络支持很多宽度那么训练时不可能每个 iteration 都遍历所有宽度。比如从 (0.25\times) 到 (1.0\times)中间可能有大量可选宽度。如果全部训练成本会非常高。所以 US-Net 必须解决一个问题每次训练只采样少数几个宽度如何让整个宽度范围都学好论文提出了一个非常简单但有效的训练规则Sandwich Rule。它的做法是每个 mini-batch 不随机采样所有宽度而是固定包含三类宽度最大宽度。最小宽度。若干随机中间宽度。例如每次训练 4 个宽度那么可以包含、再随机采样两个中间宽度。为什么这样有效因为最大宽度通常代表模型能力上界最小宽度代表模型能力下界。只要最大和最小宽度都被稳定优化中间宽度通常也能被更好地约束。所以 Sandwich Rule 的本质是用最小宽度和最大宽度夹住整个宽度空间再用随机中间宽度覆盖细节。这个策略简单但非常重要。它避免了完全随机采样导致最小宽度训练不足也保证最大宽度始终保持较强表达能力。七、关键问题三小宽度子网络怎么学得好US-Net 中所有宽度共享权重但不同宽度的容量差别很大。最大宽度子网络参数最多表达能力最强最小宽度子网络通道少容量小如果只用普通 hard label 训练很容易性能不足。为了解决这个问题论文提出Inplace Distillation原位蒸馏。它的核心做法是在同一个 mini-batch 中先运行最大宽度网络用真实标签训练最大宽度然后把最大宽度网络的输出作为软标签去监督其他较小宽度子网络。也就是说最大宽度子网络同时扮演 teacher小宽度子网络扮演 student。这个设计很巧妙因为它不需要额外训练一个 teacher 模型也不需要额外蒸馏阶段。teacher 和 student 都来自同一个 US-Net。传统蒸馏是先训练 teacher再训练 student。US-Net 的 inplace distillation 是同一个网络内部最大宽度实时指导小宽度。这对小宽度模型尤其重要。论文实验显示inplace distillation 对小宽度子网络有明显提升而对最大宽度影响很小。因此US-Net 的第三个核心技巧是用最大宽度网络作为内部 teacher提升所有小宽度子网络的性能。八、完整训练流程US-Net 的训练过程可以概括为以下步骤。第一确定宽度范围例如 ([0.25,1.0])。第二每个 iteration 根据 Sandwich Rule 选择若干宽度包括最小宽度、最大宽度和随机中间宽度。第三先运行最大宽度网络用真实标签计算损失。第四把最大宽度网络的输出固定下来作为其他宽度的蒸馏目标。第五运行最小宽度和随机中间宽度用最大宽度输出作为软标签计算蒸馏损失。第六累积多个宽度的损失更新共享权重。第七训练完成后对需要使用的目标宽度重新计算 BN post-statistics。第八部署时根据设备预算选择对应宽度运行。这个流程的核心逻辑是最大宽度保证上限小宽度保证下限中间宽度通过采样覆盖最大宽度同时作为 teacher帮助小宽度训练。九、实验设置论文主要在三个任务上验证 US-Net。第一个任务是ImageNet 图像分类主要使用 MobileNet v1 和 MobileNet v2。MobileNet v1 是较典型的轻量 CNNMobileNet v2 包含 inverted residual 结构因此可以验证 US-Net 在不同移动端架构上的效果。论文结果显示US-Net 在 MobileNet v1 和 MobileNet v2 上都优于单独训练模型和 4-switch slimmable baselines 的平均表现。(arXiv)第二个任务是DIV2K 图像超分辨率使用 WDSR 网络。这个实验说明 US-Net 的思想不只适用于分类网络也可以扩展到低层视觉任务。(arXiv)第三个任务是Atari Breakout 强化学习用来验证 US-Net 在强化学习场景中的可行性。论文摘要也提到除了 ImageNet 分类US-Net 还在图像超分辨率和深度强化学习任务上做了实验。(arXiv)这说明作者并不是只想证明一个 ImageNet 技巧而是希望说明宽度可伸缩训练是一种更通用的网络训练范式。十、实验结果解读10.1 ImageNet 上US-Net 优于单独训练和固定宽度 Slimmable Network在 MobileNet v1 和 MobileNet v2 上US-Net 的平均表现优于 individually trained networks 和 4-switch slimmable networks。论文摘要和 CVF 页面都明确提到US-Net 在 ImageNet 分类任务中相较于单独训练模型和 4-switch slimmable baseline 有更好的表现。(arXiv)这个结果很重要因为它说明共享权重并不一定会削弱子网络性能。直觉上人们可能认为每个宽度单独训练应该更好因为它没有权重共享冲突。但 US-Net 通过 Sandwich Rule 和 Inplace Distillation使小宽度子网络能够从大宽度网络中获得更强监督因此某些宽度甚至可以超过单独训练的小模型。这说明 US-Net 的优势不仅是节省训练成本也可能带来更好的小模型性能。10.2 任意宽度带来连续 FLOPs–Accuracy 曲线US-Net 的一个重要价值是训练一次后可以直接评估大量宽度下的 accuracy–FLOPs 曲线。传统模型压缩通常是一个预算训练一个模型。如果想知道 150M、200M、250M、300M FLOPs 下的精度就需要训练或剪枝多个模型。US-Net 不一样。它训练完成后可以直接选择不同宽度重新统计 BN然后测试精度。所以它能提供一条连续的精度–计算量曲线。这对部署非常重要因为真实设备预算往往不是整齐的 (0.25\times)、(0.5\times) 这种离散点而可能是任意延迟或功耗约束。这里的核心意义是US-Net 让模型部署从“选择一个固定模型”变成“在同一个模型内选择合适宽度”。10.3 Sandwich Rule 稳定了整个宽度区间论文的消融实验显示单纯随机采样宽度不如 Sandwich Rule 稳定。原因很直观随机采样可能导致最小宽度训练不足而最小宽度往往是最难训练的子网络。Sandwich Rule 每次都训练最小和最大宽度使得两个端点都被充分优化。中间宽度处在这两个端点之间因此也更容易获得稳定性能。可以这样理解最大宽度负责保证模型能力上界。最小宽度负责保证压缩极限下的性能。随机中间宽度负责覆盖连续范围。这个设计虽然简单但对 US-Net 训练非常关键。10.4 Inplace Distillation 明显提升小宽度网络小宽度子网络容量有限如果直接使用真实标签训练监督信号比较硬容易学得不够好。Inplace Distillation 让小宽度网络学习最大宽度网络的 soft prediction其中包含类别之间的相似关系比 one-hot label 信息更丰富。例如图像中某个类别虽然真实标签是 “cat”但最大宽度网络可能同时给 “tiger”“dog”“fox” 一些非零概率。这种软分布可以帮助小模型学习更细致的类别关系。所以Inplace Distillation 的作用不是简单提高训练技巧而是让小宽度子网络继承最大宽度子网络的知识。这也是 US-Net 能让小宽度表现较好的关键原因。10.5 BN post-statistics 让任意宽度真正可用没有 BN post-statisticsUS-Net 很难真正支持任意宽度。因为不同宽度下BN 的均值和方差不同。如果使用错误 BN 统计量即使权重本身训练好了推理精度也会明显下降。BN post-statistics 的意义在于宽度切换后只需要用少量样本重新校准 BN而不需要重新训练权重。这使 US-Net 的任意宽度运行变得实际可行。十一、方法优点11.1 一个模型适配多个资源预算US-Net 最大的优点是只训练和保存一个模型就可以覆盖多个计算预算。这在实际部署中非常有价值。不同设备可以选择不同宽度不需要为每个设备单独训练、保存和维护一个模型。11.2 支持任意宽度而不是固定几个宽度早期 slimmable network 只能在预定义宽度之间切换。US-Net 支持一个范围内的大量宽度选择。这使模型可以更精细地适配延迟、功耗、内存或 FLOPs 约束。11.3 不需要额外 teacherInplace Distillation 直接使用最大宽度子网络作为 teacher。这比传统知识蒸馏更方便因为不需要额外训练 teacher也不需要保存额外模型。11.4 训练技巧简单有效Sandwich Rule 和 Inplace Distillation 都很简单容易实现但对训练效果影响明显。这也是这篇论文的重要价值它不是依赖复杂结构而是通过训练策略提升共享权重子网络的性能。11.5 为后续弹性网络和一次训练多模型提供基础US-Net 的思想后来影响了很多 elastic network、once-for-all network 和 one-shot NAS 方法。比如 AutoSlim 就利用 slimmable network 来快速评估不同通道配置并进一步搜索每层通道数。AutoSlim 论文也明确将 slimmable model 用作一次性评估不同 channel configurations 的基础。(arXiv)因此US-Net 的意义不仅在于一个具体模型而在于提出了一种共享权重、多宽度、一次训练、多预算部署的训练范式。十二、方法局限12.1 它不是固定小模型US-Net 本身不会输出一个固定压缩模型而是保留完整网络并根据宽度选择子网络运行。如果部署目标是得到一个完全固定、无需宽度切换、无需 BN 校准的小模型那么 US-Net 不是最直接的方案。12.2 通道选择是前缀式的US-Net 通常使用前 (k) 个通道构成窄网络。这意味着通道存在隐式顺序前面的通道会被所有宽度共享后面的通道只在大宽度下使用。这种方式简单高效但也有限制它不能为每个宽度自由选择任意通道组合。例如某个宽度真正最优的通道集合可能不是前 (k) 个通道但 US-Net 默认只能采用前缀子网络。12.3 需要重新统计 BN虽然 BN post-statistics 成本不高但它仍然是额外步骤。如果部署时频繁切换到很多从未校准过的新宽度就需要提前准备这些宽度的 BN 统计量或者在部署端进行校准。这在某些严格部署环境中可能不方便。12.4 共享权重存在优化冲突所有宽度共享一套权重这会带来不同子网络之间的训练冲突。大宽度希望后面通道充分发挥作用小宽度则希望前面少数通道足够强。不同宽度对同一组权重的梯度方向可能并不完全一致。Sandwich Rule 和 Inplace Distillation 可以缓解这个问题但不能完全消除共享权重带来的限制。12.5 主要是统一宽度倍率US-Net 主要讨论的是所有层使用同一个 width multiplier。但实际网络中不同层的冗余程度并不相同。有些层可以剪得更多有些层对精度更敏感。统一宽度虽然简单但不一定是最优结构。后续一些工作会进一步研究每层独立通道数搜索这也是 US-Net 思想继续发展的方向。十三、整体评价这篇论文的核心价值在于它把模型压缩从“为每个预算训练一个模型”推进到“一个模型覆盖多个预算”。US-Net 不直接删除通道也不根据每张输入动态选择路径而是根据部署预算选择宽度。它更关注的是budget-adaptive inference也就是给定一个资源预算后模型可以立刻切换到合适宽度。它的三个关键技术分别解决了三个核心问题BN post-statistics 解决任意宽度下 BN 统计不一致的问题。Sandwich Rule 解决不能训练所有宽度的问题。Inplace Distillation 解决小宽度子网络性能不足的问题。这三个技巧组合起来使一个共享权重网络可以在很宽的宽度范围内稳定运行。十四、一句话总结《Universally Slimmable Networks and Improved Training Techniques》提出 US-Net使一个神经网络能够在给定范围内以任意宽度运行它通过 BN post-statistics 解决不同宽度的 BN 统计问题通过 Sandwich Rule 高效训练最小、最大和随机中间宽度并通过 Inplace Distillation 让最大宽度子网络指导小宽度子网络从而实现一次训练、多宽度部署为后续弹性网络、AutoSlim 和 Once-for-All 类方法奠定了重要基础。