应对未来物联网大潮：如何在内存有限的情况下部署深度学习？

对于任何想要构建大型服务的人来说，部署需要大量内存的深度学习算法都是一项挑战。从长期使用的角度来说，云服务过于昂贵。相比之下，在边缘设备上离线部署模型更为便宜，并且还有许多其他好处。这种做法唯一的缺点就是，边缘设备缺乏内存和计算能力。

本文探讨了一些可用于在内存受限的配置下部署神经网络的技术。由于训练和推理阶段所使用的技术不同，这两部分将分开来讲。

训练阶段

某些应用必须要在线学习。也就是说，模型性能的提高是依赖于反馈或附加数据的。在边缘设备上部署此类应用程序会对模型造成切实的资源限制。以下 4 种方法可以减少此类模型的内存消耗。

1. 设置梯度检查点

TensorFlow 这类框架的训练需要消耗大量内存。在前向传递期间，计算图中每个节点的值都需要被计算并保存在内存中。这是在反向传播期间计算梯度所必需的。

图：节点中的每个值都要被保存下来，用以在一次反向传播中计算梯度。（来源：https://github.com/openai/gradient-checkpointing）

通常情况下这是可以的，但是当模型越来越复杂时，内存消耗会急剧增加。解决此问题的一个简洁的方案是在需要时重新计算节点的值，而不是将它们保存到内存中。

图：重新计算节点的值用以计算梯度。注意，我们需要好几次分步的前向传递，才能计算一次反向传播。（来源：https://github.com/openai/gradient-checkpointing）

然而，从上图可以看出，这样做会使计算成本显著增加。一个比较好的权衡是在内存中只保存一些节点，同时在需要时重新计算其他节点。这些保存的节点称为检查点。这大大减少了深度神经网络内存消耗。如下图所示：

图：左数第二个节点为检查点。这种做法可以减少内存消耗，为此增加的计算时间也在接受范围之内。（来源：https://github.com/openai/gradient-checkpointing）

2.牺牲速度换取内存（重计算）

将上述想法进一步扩展，我们可以重新计算某些操作以节省内存的消耗。内存高效的 DenseNet 的实现就是一个很好的例子。

图：DenseNet 中的一个稠密块。（来源：https://arxiv.org/abs/1608.06993）

DenseNet 中的参数效率很高，但其内存效率却很低。产生这种矛盾的原因是 DenseNet 的拼接结构及 batchnorm 本身的性质。

想要在 GPU 上实现高效卷积，数值必须连续放置。因此，在拼接操作之后，cuDNN 会将值在 GPU 上连续排列。这会产生大量的冗余内存分配。同样，batchnorm 也会产生过多的内存分配，如论文中所述。这两种操作都会导致内存呈平方级增长。DenseNet 结构中含有大量的拼接操作和 batchnorm，因此其内存效率十分低下。

图：直接的拼接和 batchnorm 操作及其高效内存实现的对比。（来源：https://arxiv.org/pdf/1707.06990.pdf）

针对上述问题，有一个简洁的解决方案。该方案基于两个关键的现象。

首先，拼接和 batchnorm 操作不是时间密集的。因此，我们可以在需要时重新计算其值，而不是存储所有冗余的内存。其次，我们可以使用“共享内存空间”来转储输出，而不是为输出分配“新”内存空间。

这个共享空间可以被覆盖，用来存储其他拼接操作的输出。我们可以在需要时对拼接操作进行重新计算，用以计算梯度。同样，我们可以将这种做法扩展到 batchnorm 操作。这个简单的技巧可以节省大量的 GPU 内存，而计算时间不会增加太多。

3.牺牲浮点数精度

在一篇优秀的博客（链接：https://petewarden.com/2015/05/23/why-are-eight-bits-enough-for-deep-neural-networks/）中，Pete Warden 解释了如何使用 8 比特浮点数训练神经网络。但由于浮点数精度的降低，会产生许多问题，其中一些问题如下：

• 如这篇论文中所述（链接：https://arxiv.org/pdf/1412.7024.pdf） ，“激活函数、梯度和参数”的取值范围大不相同。用一个固定点来表示的方式不会很理想。论文表示，可以设置一个“动态固定点”表示，这对于低浮点数精度的神经网络会很有效。

• Pete Warden 在其另一个博客（链接：https://petewarden.com/2017/06/22/what-ive-learned-about-neural-network-quantization/） 中讲到，浮点数精度越低意味着预测值距正确值的偏差越大。一般来讲，如果错误完全随机，那么他们相互之间很可能会抵消。然而，在 padding，dropout 和 ReLU 操作中，0 这个值被广泛采用。而如果采用低浮点数精度格式，0 这个值很难被精确地表示，因此会带偏整个网络的性能。

4.调整网络架构

调整网络架构也就是设计一个新的神经网络结构，该结构需要做到优化精度、内存和速度。下面讲述几种可以在空间或时间上优化卷积操作的方法。

• 将 NxN 的卷积操作分解为 Nx1 和 1xN 的卷积操作的组合。这样会节省大量的空间，同时也会提高运算速度。这个小窍门和其他的一些方法已经被用在了新版本的 Inception 网络中。如果想要详细了解，请参见这个博客：https://towardsdatascience.com/a-simple-guide-to-the-versions-of-the-inception-network-7fc52b863202。
  

• 使用深度可分离的卷积结构（Depthwise Separable convolutions）。MobileNet（详见 https://arxiv.org/pdf/1704.04861.pdf）和 XceptionNet（详见 https://arxiv.org/abs/1610.02357）中就使用了该结构。关于这种卷积结构的详细讨论，参见这个博客：https://towardsdatascience.com/types-of-convolutions-in-deep-learning-717013397f4d。
  

• 使用 1x1 的卷积作为瓶颈层，以减少特征的通道数。这种技术被用在了好几个较为流行的神经网络中。

图：谷歌自动化机器学习 Google AutoML 示意图。（来源：https://www.youtube.com/watch?v=Y2VF8tmLFHw）

关于这一问题，还有一个非常有趣的解决方案，即针对具体问题让机器自己决定哪个架构最好。网络结构搜索技术（链接：https://arxiv.org/pdf/1707.07012.pdf）利用机器学习为给定的分类问题寻找最好的神经网络结构。在ImageNet上，使用该技术生成的网络（NASNet）是目前性能最好的。谷歌的自动化机器学习（链接：https://ai.googleblog.com/2017/11/automl-for-large-scale-image.html）就是基于这一理论构建的。

推理阶段

在推理阶段将模型部署到边缘设备相对来说比较简单。本节介绍了几种技术，可以为这种边缘设备优化你的神经网络。

1.去掉“冗余部件”

TensorFlow 这类机器学习框架往往需要消耗巨大的存储空间用于创建计算图。这些额外的空间消耗可以加速训练，但在推理阶段却没什么用。因此，专门用于训练的这部分图可以直接被剪掉。我们可以将这部分图称为“冗余部件”。

对于 TensorFlow，建议将模型的检查点转换为冻结的推理图。这一过程会自动删除占内存的冗余部件。直接使用模型检查点的图时报出的资源耗尽错误有时可以通过这种转换为冻结推理图的方式解决。

2.特征剪枝

Scikit-Learn 上的一些机器学习模型（如随机森林和 XGBoost）会输出名为 feature_importances_的属性。这一属性代表了在当前分类或回归问题中每个特征的显著性。我们可以直接将不显著的特征剪枝掉。如果你的模型中的特征非常非常多，又难以通过其他方法降低特征量，那么这种方法会非常有效。

图：特征重要性图示例。（来源：https://machinelearningmastery.com/feature-importance-and-feature-selection-with-xgboost-in-python/）

相似地，在神经网络中，很多权重的值很接近于 0。我们可以直接剪去那些权重趋于 0 的连接。然而，移除层间独立的连接可能会产生稀疏矩阵。不过，已经出现了高效推理器（硬件）的相关工作（链接：https://arxiv.org/pdf/1602.01528.pdf），可以无缝处理稀疏矩阵问题。然而，大部分的机器学习框架还是将稀疏矩阵转换为稠密格式，然后传输到GPU上。

图：删除不显著的卷积核。（来源：http://machinethink.net/blog/compressing-deep-neural-nets/）

另一种做法是，我们可以移除不显著的神经元，并对模型稍加重新训练。对于 CNN 来说，我们也可以移除整个卷积核。研究（链接：https://arxiv.org/pdf/1510.00149.pdf）和实验（链接：http://machinethink.net/blog/compressing-deep-neural-nets/）表明，使用这种方法，我们可以做到在最大程度上保持精度，同时减少模型的大部分空间占用。

3.共享权重

为更好地解释权重共享，我们可以考虑一篇关于深度压缩（链接：https://arxiv.org/pdf/1510.00149.pdf）的论文中的例子。考虑一个4x4的权重矩阵。它含有16个32比特浮点值。

我们可以将权重值量化为 4 个层，但保留其 32 比特的属性。现在，这个 4x4 的权重矩阵只有 4 个值了。这 4 个不同值被保存到一个独立的（共享的）内存空间中。我们可以分别给这 4 个不同值赋予一个 2 比特的地址（比如 0，1，2，3）。

图：权重共享示意。（来源：https://arxiv.org/pdf/1510.00149.pdf）

我们可以使用 2 比特地址来索引权重值。因此，我们获得了一个由 2 比特地址表示的新的 4x4 矩阵，矩阵中的每个位置代表其在共享存储空间中的位置。使用该方法表示整个矩阵仅需要 160 比特（16 * 2 + 4 * 32）。这样，我们获得了 3.2 的尺寸缩减系数。

当然，这种尺寸的缩减所伴随的将是时间复杂度的增加。但是，访问共享内存的时间不会成为非常严重的时间损失。

4.量子化和降低浮点数精度（推理阶段）

回想一下我们在前文中训练部分讲到的降低浮点数精度的方法。对于推理阶段，降低浮点数精度不会像训练阶段那样麻烦。权重可以直接被转换为低精度格式（详见https://heartbeat.fritz.ai/8-bit-quantization-and-tensorflow-lite-speeding-up-mobile-inference-with-low-precision-a882dfcafbbd），并直接用于推理。不过，如果想要将精度降低很多，那么权重需要稍作调整。

5.编码

通过使用编码，可以进一步对修剪和量化过的权重的空间占用进行优化。霍夫曼编码可以用较低的位数表示使用最频繁的权重值。因此，在比特级别，经过霍夫曼编码的字符串比正常的字符串的空间占用更少。

深度压缩使用的是无损压缩技术（如霍夫曼算法）进行压缩。而也有研究使用了有损压缩技术（详见https://arxiv.org/abs/1711.04686）。这两种方法都有一个缺点，即翻译过程的开销过大。

6.推理优化器

到目前为止，我们已经讨论了一些很棒的想法，但是从头开始实施它们需要相当长的时间。在这里，推理优化器就有了用武之地。例如，Nvidia 的 TensorRT 结合了所有这些伟大的想法（甚至还包括更多），并在训练好的神经网络的基础上提供了优化的推理引擎。

图：TensorRT。（来源：https://developer.nvidia.com/tensorrt）

此外，TensorRT 还可以优化模型，使其更好地利用 Nvidia 的硬件。下面是一个示例，其中使用 TensorRT 优化过的模型能更有效地兼容 Nvidia V100。

图：在 Nvidia V100 上使用 TensorRT 优化过的模型。（链接：https://devblogs.nvidia.com/tensorrt-3-faster-tensorflow-inference/）

7.知识蒸馏

我们可以教小型模型来模仿强大的大型模型的性能，而不是执行花哨的优化技术。这种技术叫做为知识蒸馏，它被集成在 Google Learn2Compress 中。

图：教师-学生模型。（来源：https://ai.googleblog.com/2018/05/custom-on-device-ml-models.html）

通过使用此方法，我们可以强行使边缘设备上的较小模型达到较大模型的性能水平。研究表明，使用此方法所造成的精度下降很小。如果想要详细了解此技术，请参考 Hinton 的论文（链接：https://arxiv.org/pdf/1503.02531.pdf）。

英文原文链接：https://heartbeat.fritz.ai/how-to-fit-large-neural-networks-on-the-edge-eb621cdbb33?mkt_tok=eyJpIjoiTmpVMFpURTBOakV6TXpOaiIsInQiOiJjRE5YUmJ3dmlFcEpRMUZ0bFlNUmJjRmtsWUxZRGowVkJlSnZhNUJvUEhmeGs0bWhnV1JkNnh6XC9zRU9ja3BKeWowUjlNWjFaSDhQWlNwTHgxYm9IY1BqWEh0V1lrMGxKYjlQMDZId25aRWZyZ0Z3R1ZpYjhISFhlVVcwOWh3ODUifQ%3D%3D