每天一天论文 364/365 Memory-Efficient Implementation of DenseNets

Memory-Efficient Implementation of DenseNets

Code

Densely Connected Convolutional Networks (DenseNets)

摘要

由于特性重用，DenseNet体系结构[9]具有很高的计算效率。然而，一个简单的DenseNet实现可能需要大量的GPU内存:如果管理不当，预激活批处理规范化和连续卷积操作可以生成随网络深度二次增长的特征图。在本文中，我们介绍了在训练中减少DenseNets内存消耗的策略。通过策略性地使用共享内存分配，我们将存储特征映射的内存成本从二次型降低到线性型。没有GPU内存瓶颈，现在可以训练极深的densenet。拥有14M参数的网络可以在单个GPU上训练，而原来是4M。一个264层的DenseNet (73M参数)，以前是不可行的培训，现在可以在一个工作站与8英伟达特斯拉M40 gpu培训。在ImageNet ILSVRC分类数据集上，这个大型DenseNet获得了最先进的单作物top-1错误，为20.26%。

贡献

本文提出一种策略，该策略可以显著降低DenseNet实现的训练时间内存成本，同时略微降低速度。我们的主要观察是，负责大部分内存消耗的中间特征映射的计算成本相对较低。这允许我们引入共享内存分配，所有层都使用它来存储中间结果。后续层会覆盖前一层的中间结果，但是可以在反向传递期间以最小的代价重新填充它们的值。这样做可以减少二次型到线性型的feature map内存消耗，同时只增加15 - 20%的额外训练时间。这种内存节省使得在合理的GPU预算下训练超大的densenet成为可能。特别是，我们能够将最大的DenseNet从161层(每层有48个特征，20M参数)扩展到264层(每层有48个特征，73M参数)。在ImageNet上，这个模型实现了20.26%的单作物top-1错误，这(就我们所知)是最先进的。

1. 预激活Batch normalization
   
   DenseNet体系结构利用预激活批处理规范化[7]。与传统的结构不同，预激活网络在卷积运算之前而不是之后应用批处理标准化和非线性。虽然这可能看起来是一个小的变化，但它在DenseNet性能上有很大的不同。批处理标准化对输入特性应用扩展和偏差。如果每一层都应用它自己的批处理标准化操作，那么每一层都对以前的特性应用唯一的比例和偏差。例如，第2层批处理规范化可能将第1层的特性伸缩一个正的常数，而第3层可能将相同的特性伸缩一个负的常数。在非线性之后，第2层和第3层从第1层提取相反的信息。如果所有层共享同一批归一化操作，或者在卷积操作之后进行归一化，则不可能实现。如果不进行预激活，100层的DenseNet-BC的CIFAR-10错误将从4.51增加到5.18。在CIFAR-100上，错误从22.27增加到24.30(图2左)。
2. Naïve Implementation and Efficient implementaion
   每一层的输入都是前一层的特征(有颜色的盒子)。由于这些特性来自不同的层，所以它们在内存中不是连续存储的。因此，Naïve实现的第一个操作复制这些特性，并将它们连接到一个连续的内存块(中间偏左)。如果前面的每一层都产生k个特征，那么相邻的内存块必须容纳k个特征映射。这些连接起来的特性被输入到批处理规范化操作(右中)，该操作类似地分配一个新的’×k连续的内存块。(ReLU非线性发生在内存中，因此为了简单起见，我们选择将其排除在计算图之外。)最后，卷积运算(最右边)从批量归一化输出生成k个新特征。从图3(左上角)可以很容易地看到内存的二次增长。这两个中间操作需要一个内存块，它可以适合所有O(’ k)先前计算的特性。相比之下，输出特性每层只需要一个常数O(k)内存。
   
   图3 DenseNet层向前传递:原始实现(左)和高效实现(右)。实体盒对应于内存中分配的张量，其中作为透明盒的是指针。实箭头表示计算，虚线箭头表示内存指针。高效的实现将连接层、批处理规范化层和ReLU层的输出存储在临时存储缓冲区中，而原始实现分配新的内存。

为了避免这个问题，我们利用连接和标准化操作在计算上是非常便宜的。为了避免二次内存增长，我们提出了两个预先分配的共享内存存储位置。在转发期间，我们将所有中间输出分配给这些内存块。在反向传播期间，我们根据需要动态地重新计算连接和标准化的特性。

Shared storage for concatenation：第一个操作—特性连接—需要O(’ k ')内存存储，其中’是前一层的数量，k是每一层添加的特性的数量**。我们不是为每个连接操作分配内存，而是将输出分配给跨所有层共享的内存分配(图3中的“共享内存存储1”)**。因此，连接特性映射(中间偏左)是指向这个共享内存的指针。复制到预先分配的内存要比分配新内存快得多，因此这种连接操作非常有效。批处理规范化操作将这些连接的特性作为输入，直接从共享内存存储1读取数据。因为所有网络层都使用共享内存存储1，所以它的数据不是永久的。当在反向传播过程中需要连接的特征时，我们假设可以有效地重新计算它们。
Shared storage for batch normalization：类似地，我们将批处理规范化(也需要O(m)内存)的输出分配给共享内存分配(“共享内存存储2”)。卷积操作从这个共享存储的指针中读取数据(中右)。与连接一样，我们必须在反向传播期间重新计算批处理规范化输出，因为共享内存存储2中的数据不是永久的，将被下一层覆盖。然而，批处理归一化由标量乘法和加法运算组成，与卷积运算相比效率非常高。计算批处理规范化特性大约占一次正反向传递的5%，因此重复操作成本不高。
**Shared storage for gradients:**在反向传播过程中，连接、批处理归一化和卷积操作都会产生梯度张量。在LuaTorch中，确保这个张量数据存储在一个共享内存分配中是很简单的。这种策略防止了梯度存储的二次增长。我们使用基于Facebook的ResNet实现的内存共享方案。PyTorch和MxNet框架共享开箱即用的梯度存储
**Putting these pieces together:**前向传递(图3右)类似于navi实现(图3左上)，除了中间功能映射存储在共享内存存储1或共享内存存储2中。向后传递需要一个额外的步骤:我们首先重新计算连接和批处理规范化操作(中间左和中间右)，以便用适当的feature map数据重新填充共享内存存储。一旦共享内存存储包含正确的数据，我们就可以执行常规的反向传播来计算梯度。总的来说，这个DenseNet实现只为输出特性分配内存(最右边)，输出特性的大小是不变的。它对地形图的总内存消耗在网络深度上是线性的。

结论

内存消耗
这三种实现(在LuaTorch和PyTorch中)的内存消耗如图4所示。(没有PyTorch的简单实现，因为PyTorch自动共享梯度存储。)有了四个二次运算，简单的实现很快就变成了内存密集型的。一个160层的网络(k = 12个特性/层，1.8M参数)的内存使用量大约是一个40层网络(k = 12, 160K参数)的10倍。训练一个超过160层的更大的网络需要超过12 GB的内存，这比一个典型的GPU要多。虽然共享梯度减少了一些内存开销，但是内存消耗仍然随着深度快速增长。另一方面，使用所有内存共享操作可以显著减少内存消耗。在LuaTorch中，160层模型使用了初始实现所需内存的22%。在相同的内存预算(12gb)下，可以训练一个340层的模型，它的深度是最佳初始实现模型的2.5倍，参数是最佳初始实现模型的6倍。

在图4中，我们可以看到PyTorch比LuaTorch具有更高的内存效率。使用高效的PyTorch实现，我们可以在一个GPU上训练近500个层(13M参数)的densenet。PyTorch中的“autograd”库在训练期间执行内存优化，这可能有助于提高这个实现的效率。

训练时间
训练时间不会受到内存优化的显著影响。在图5中，我们绘制了100层DenseNet-BC (k = 12)在NVIDIA Maxwell Titan-X上的每一小批时间。共享梯度存储不会产生任何时间成本。共享批处理规范化和连接存储在LuaTorch上大约增加15%的时间开销，在PyTorch上增加20%的时间开销。
误差分析
我们在图6中显示了这些densenet分类误差性能。使用有效实现训练的模型表示为绿点(表示标准学习速率调度的平方，表示余弦的星星)。我们将这些模型的性能与用原始实现训练的浅层次DenseNets进行了比较。此外，我们还比较了[6]中引入的ResNet模型和[14]中引入的ResNeXt模型。对单中心试验图像进行了分析