[TinyML]NetAug:Network Augmentation for Tiny Deep Learning

NetAug:基于TinyDL的网络增强

from ICLR2022


MIT HAN Lab

Abstract

本文提出了网络增强（NetAug)这种新的训练方法来提升微型网络的性能。现有的正则化方法（数据增强、随机失活）等已经在大型神经网络上取得了成功，一般通过添加噪声来克服过拟合。但是我们发现这些方法会降低微型网络的性能、本文认为：小型网络的训练不同于大型网络，不应该增强数据而是应该增强网络。因为小型网络一般容量有限，一般都是欠拟合而不是过拟合。

NetAug增强了网络（通过反随机失活），而不是在数据集或网络中插入噪声。它将小模型插在更大的模型中并且鼓励其作为更大模型的子模型，这样可以获得额外的监督。在测试阶段只有小型模型进行推理，这样是零推理开销。本文将NetAug在图像分类和目标检测任务上进行了测试，在ImageNet上提升了2.2%的分类精度，在COCO数据集上MAX减少38%在PascalVOC数据集上MAC减少41%。

Section I Introduction

TinyIoT设备正在快速增长，直到2025年将达到754.4亿台，在这些微型边缘设备上直接步数深度神经网络而不需要云端服务器的参与将会带来更好的隐私性，同时降低系统成本。但是与云端设备相比（GPU）边缘设备的资源十分有限。比如一个STM32 MCU只有320kB的内存比GPU小50000倍。因此在上面部署的神经网络必须非常小才能在边缘设备上运行。



因此提升微型神经网络的性能已经成为微型深度学习的基本挑战。
传统方法依赖正则化、随机失活来减轻过拟合，比如数据增强（自动增强、Mixup)，但是对于小型网络并不适用。参见Fig 1可以看到左侧ResNet-50使用正则化后可以有效提升性能，但是对MobieNetV2-Tiny这种比ResNet-50小174倍的网络确实有害的。本文认为小型网络的训练与大型网络有本质的不同，不应该扩充数据及而是扩充网络。大型网路哦倾向于对训练数据进行过拟合，数据增强可以有效缓解但是微星网络由于容量有限，对数据的学习是够的，因此再使用正则化则会进一步降低性能。

本文提出了网络增强（NetAug),本文认为小型网路偶在训练时需要更大的容量而不是噪声。因此NetAug并没有增强数据，也没有使用随机失活，而是将小模型插入到更大的模型中，通过权重和梯度共享来增强小型网络，此外小模型还会作为大模型的子模型，可以看做是dropout的逆变换，因为本文扩大了目标模型而不是缩小目标模型。训练时NetAug会加入大模型的梯度作为对小模型的额外监督；测试时则只使用小模型进行推理，从而是zero-inference.

Section II Related Work

Knowledge distillation

知识蒸馏会将从教师模型中学到的dark knowledge迁移到小的学生模型上，训练学生模型的输出来匹配教师模型的输出。KD经常和其他方法结合来共同提升网络性能，如逐层KD，网络剪枝。

与KD不同之处在于本文解决的是小型网络欠拟合的问题。因此技术实现上并不需要训练目标模型来模拟教师模型，而会将目标模型作为一个更大网络的子网络，通过增加目标网络的宽度来获得额外监督。由于这种方法与KD有根本上的不同因此可以结合起来联合提升模型性能。

Regularization Methods

正则化方法可以分为数据增强和随机失活两类。数据增强会向数据集中添加噪声，如MixUp,Cutout。还有借助AutoML的如AutoAugment,RandAugment
随机失活则是通过向网络中增加噪声，比如Dropout会随机失活网络连接；还有StochasticDepth,SpatialDropout.DropBlock等。还有将KD与其他方法相结合的：Self-distillation ，GradAug.
本文与上述正则化的不同之处在于通过增加网络的宽度而不是通过随机下降来提升小型网络的性能，可以看做是一种reversed Drouput的形式，在小型网络上比正则化更有效。

Tiny Deep Learning

在Tiny DL方向，一种是压缩和量化现有神经网络，另一种则是高效神经网络设计，此外除了手动设计压缩策略和网络架构还有借助AutoML来自动模型压缩和自动网路设计。本文的方法侧重提升小型网络的精度，不改变网络的结构，在实验中取得了更好的性能。

Section III Network Augmentation

Part 1 Formulation

将小网络的权重表示为Wt，损失函数为L，则小网络的训练就是不断最小化损失函数，使用随机梯度下降，但由于小网络容量有限，所以性能一般较差，会落入局部最优。
因此本文在训练过程中引入额外的监督信息，与随机失活相反，NetAug将小网络作为大网络的一个子模型，则损失函数表示为：

每一对代表小网络的权重Wi,alpha代表不同加强模型的权重.

从Fig 2可以看出 本文通过将小网络放入更大啊网络中来增强整个网络，二者权重是共享的，这一tiny network作为更大网络的监督向，每次训练都会采样一个tiny network提供额外的监督信息，测试的时候只采用tiny network做推理。因此可以将NetAug理解为通过增加tiny network的宽度和扩展比实现的。

Part 2 Implementation

搭建增强网络：

将每一个增强模型的权重都保存下来不现实，因为网络模型会随着增强模型的数目线性增长。因此本文在不同增强模型之间权重共享只保留最大的增强模型的权重，通过从最大的增强模型中采样来获得其他的增强强模型，参见Fig 2左侧。
这种权重共享策略也在one shot NAS和多任务学习中使用。我们的目标则与他们不同：

（1）one shot NAS会训练一个超网通过权重共享生成所有子网，从而提升NAS预测的效率。而NetAug聚焦于通过增强模型提供的监督信息来提升tiny network的性能，吧并且也可以应用到NAS中来提升模型的性能。

（2）多任务学习致力于通过权重共享迁移多个任务学到的知识，而NetAug侧重点是传输单个任务的监督信息。
通过不断增加宽度获得最大的增强模型，这样在GPU上比搭建更深层次的网络训练时间要少。比如卷积操作的宽度是w本文会乘上一个增强因子r最终的宽度就是二者的乘积。
搭建最大的增强模型后会通过选择通道子集 方式获得其他增强网络，比如宽度可以是w,2w,3w在r=3的情况下。


Training Process



可以看到每次获得增强网络的监督信息都需要一次前向和反向传播过程，计算成本十分高昂，因此本文每一步值采样一个增强网络，通过将基础的监督信息和额外的监督信息融合来更新tinynetwork的权重。
并且本文将所有增强网络的权重因子设置为1，其实每次采样也可以采样不止一个网络但是本文的实验发现不但增加了训练成本还会影响网络性能，因此本文在每一步训练时只采样一个网络。

Part 3 Training and Inference overhead

NetAug仅在训练时使用，推理的时候还是使用tiny network，因此NetAug的推理成本是0，并且由于NetAug并没有改变网络结构也不需要额外的软硬件支持，因此实际应用时更简易。
接下来再考虑NetAug的训练成本，虽然看起来每次训练时增加了一个额外的前向和反向传播的成本，但是实验发现训练时间仅仅增加了16.7%远远没有double训练时间。主要是因为tiny network的主要训练时间都在data loading和通信上了，并不是网络的前向和反向传播上。
此外NetAug也会增加网络的峰值内存占用，因为tiny network的峰值内存占用要小的多但是实践中借助GPU训练这部分的内存增加是可以接受的。

Section IV Experiments

数据集

 ImageNet, Food101，Flowers102,Cars,Cub200,Pets,PASCAL VOC,COCO

 
主要有分类、细粒度分类、目标检测

Part 1 Results on ImageNet

本文与MobineNetV2-Tiny,ProxylessNAS,MCUNet等小网络进行了对比，结果参见Table 1.

 



可以看到使用NetAug都对基线网络有进一步的提升，并且对越小的网络提升越明显。我们认为小网络的容量更小，因此欠拟合问题更显著也能从NetAug中获益更多，因此NetAug影响了Resnet50的精度也就不奇怪了。
从Fig 3也可以看到NetAug会一直提升网络精度，或者在达到相同精度的前提下NetAug的训练成本更低：减少了ImageNet上75%的训练时间，或者IageNet21K 83%的训练epoch。

Comparison with KD


本文还对比了做知识迁移的性能，比较结果参见Table 2.将NetAug与KD结合可以共同提升tiny network的分类性能：




Comparison with Regularison methods



由于tiny network的容量十分有限因此即使十分微小的正则化也会损伤网络性能，像强度更大的RandAugment,Dropblock等会更加影响网络的精度。本文还测试了正则化的强度对网络性能的影响，参见Table 6,有以下发现：增加正则化会影响tiny network的性能，并且正则化强度越大网络精度下降的越大。
基于上述观察本文认为tiny network面临的更多是欠拟合而非过拟合的问题，应用正则化并不足以解决这一问题，而应该使用netAug会将MobileNetV2的分类精度提升1.3%，训练成本仅增加了16.7%。



Discussion



NetAug通过减轻欠拟合问题提升了tiny network的网络性能，对于大型网络不存在欠拟合问题那么NetAug的提升可能没那么明显，本文将验证结果展示在Fig 4中。

Part 2 Results on Transfer Learning

许多在ImageNet上预训练的模型会迁移到下游任务中，比如继续做细粒度图像分类或目标检测。因此本文继续探究了NetAug是否有助于迁移学习。
做细粒度分类的展示在Table 4，做目标检测的展示在Table 5.


我们发现在训练前的准确性更高
数据集(在本文中是ImageNet)并不总是在下游导致更高的性能。例如，虽然添加KD提高了MobileNetV2w 0.35的ImageNet精度
和MobileNetV3w 0.35，使用与KD预先训练的权重降低了在两个细粒度上的表现
粒度分类离子数据集(Cub200和Pets)和目标检测数据集(PascalVOC和
COCO)。同样地，更多epoch的训练模型显著地提高了
MobileNetV2w0.35的精度，但损害了三个细粒度分类数据集的性能。
与KD，训练更多epoch相比本文发现结合NetAug的预训练模型更适合不同场景下的迁移。这说明NetAug不仅能提升模型精度还能改善模型的表征质量。

除了提升网络性能，NetAug还能提升推理效率，Fig 5展示了输入不同分辨率的结果，可以看到在达到相近目标检测精度的前提下，NetAug输入的分辨率更低，减少了41%的MAC在PASCAL数据集上，以及减少38%MAC在COCO数据集上。更小的输入分辨率也会减少推理时 的内存占用，这在内存受限的设备上部署网络十分重要。

Section V Conclusion

本文提出NetAug来提升tint network的网络性能，解决了小网络欠拟合的问题。不同于正则化，NetAug通过将目标tiny网络插入到更大的额网络中训练获得额外的监督信息有效弥补了小网络容量有限的问题。在图像分类和目标检测等实验中均证明了NetAug在小型网络上的有效性。