[论文解读] Profile-guided memory optimization for deep neural networks

Taro Sekiyama, Takashi Imamichi+2 more|arXiv (Cornell University)|Apr 26, 2018

Neural Networks and Applications被引用 9

一句话总结

本文提出一种基于配置文件的深度神经网络（DNN）内存优化方法，通过从单次前向传播中获取的配置文件数据，采用快速启发式内存分配策略，将峰值内存使用量最高降低49.5%。该方法通过实现更快的内存分配和更大的小批量尺寸，使训练和推理速度最高提升四倍，且在配置文件分析后无运行时性能开销。

ABSTRACT

Recent years have seen deep neural networks (DNNs) becoming wider and deeper to achieve better performance in many applications of AI. Such DNNs however require huge amounts of memory to store weights and intermediate results (e.g., activations, feature maps, etc.) in propagation. This requirement makes it difficult to run the DNNs on devices with limited, hard-to-extend memory, degrades the running time performance, and restricts the design of network models. We address this challenge by developing a novel profile-guided memory optimization to efficiently and quickly allocate memory blocks during the propagation in DNNs. The optimization utilizes a simple and fast heuristic algorithm based on the two-dimensional rectangle packing problem. Experimenting with well-known neural network models, we confirm that our method not only reduces the memory consumption by up to $49.5\%$ but also accelerates training and inference by up to a factor of four thanks to the rapidity of the memory allocation and the ability to use larger mini-batch sizes.

研究动机与目标

• 为解决深度神经网络（DNN）在训练和推理过程中高内存消耗的问题，特别是在GPU和边缘系统等内存受限设备上。
• 在不改变模型计算的前提下，降低DNN的峰值内存使用量，从而支持更大的小批量尺寸并提升GPU利用率。
• 开发一种快速、基于配置文件的内存分配技术，避免深度学习框架中动态内存分配带来的性能开销。
• 在前馈网络（如CNN）和循环网络（如seq2seq）中实现高效内存管理，尤其适用于输入长度可变的场景。

提出的方法

• 通过分析DNN单次前向传播过程中的内存使用情况，捕捉‘热点’计算路径的内存访问模式。
• 将内存分配问题建模为二维矩形装箱问题，以最小化峰值内存使用量。
• 设计一种简单而快速的启发式算法，高效生成接近最优的内存块分配方案。
• 从配置文件中预计算并缓存内存地址，从而在后续运行中实现常数时间的内存访问。
• 该方法已集成至Chainer深度学习框架，并与其它框架兼容。
• 对于可变长度序列（如RNN中的输入），在必要时重新计算内存分配，以避免因未使用内存块导致的内存膨胀。

实验结果

研究问题

• RQ1 基于配置文件的内存分配能否在无运行时开销的前提下显著降低DNN的峰值内存使用量？
• RQ2 针对DNN内存消耗最小化，基于快速启发式的二维矩形装箱方法在实际中有多高效？
• RQ3 内存减少在多大程度上可支持更大的小批量尺寸并提升训练过程中的GPU利用率？
• RQ4 尽管存在配置文件计算成本，该方法是否仍能保持或提升推理速度？
• RQ5 该方法在具有可变长度输入的模型（如RNN）上表现如何？

主要发现

• 在ResNet-50和Inception-ResNet中，该方法将峰值内存使用量最高降低49.5%；在GoogLeNet和ResNet-50的推理阶段，分别实现12.6%和10.0%的内存降低。
• 由于内存分配速度加快以及小批量尺寸增大带来的更高GPU利用率，训练速度最高提升四倍。
• 在seq2seq模型中，推理速度最高提升23.8%，得益于更低的内存分配成本和更优的内存复用。
• 启发式算法执行速度快，适用于实际应用，大多数模型的执行时间均低于1秒。
• 在RNN中，该方法优于基于内存池的管理方式，避免了内存池碎片化问题，并减少了内存释放与重新分配的次数。
• 在128和256的小批量尺寸下，优化版本的训练速度超过原始Chainer，主要得益于更低的内存分配开销。