[论文解读] BrainSlug: Transparent Acceleration of Deep Learning Through Depth-First Parallelism

Nicolas Weber, Florian Schmidt+2 more|arXiv (Cornell University)|Apr 23, 2018

Advanced Neural Network Applications参考文献 17被引用 7

一句话总结

BrainSlug 通过将神经网络计算从广度优先重新排序为深度优先处理，加速了深度学习的推理与训练，提升了数据局部性并减少了缓存未命中。它在无需修改模型架构且仅需极少代码改动的情况下，实现了 CPU 上高达 41.1% 和 GPU 上高达 35.7% 的性能提升，使标准硬件上实现透明的性能增益成为可能。

ABSTRACT

Neural network frameworks such as PyTorch and TensorFlow are the workhorses of numerous machine learning applications ranging from object recognition to machine translation. While these frameworks are versatile and straightforward to use, the training of and inference in deep neural networks is resource (energy, compute, and memory) intensive. In contrast to recent works focusing on algorithmic enhancements, we introduce BrainSlug, a framework that transparently accelerates neural network workloads by changing the default layer-by-layer processing to a depth-first approach, reducing the amount of data required by the computations and thus improving the performance of the available hardware caches. BrainSlug achieves performance improvements of up to 41.1% on CPUs and 35.7% on GPUs. These optimizations come at zero cost to the user as they do not require hardware changes and only need tiny adjustments to the software.

研究动机与目标

• 解决传统广度优先处理中因数据局部性差和缓存频繁抖动导致的深度神经网络推理与训练性能瓶颈。
• 在不修改硬件或模型架构的前提下，实现对标准 CPU 和 GPU 上深度学习工作负载的透明加速。
• 开发一个模块化框架，通过可扩展的 API 支持多种深度学习框架（如 PyTorch、TensorFlow）和硬件后端（CPU、GPU、FPGA）。
• 通过以深度优先顺序处理输入数据的子集，优化计算过程，提升缓存利用率并减少冗余内存访问。
• 在多种网络架构和批量大小下评估框架性能，重点关注实际部署场景中的真实性能提升。

提出的方法

• 用深度优先遍历替代标准的逐层广度优先处理，即在处理下一个数据子集前，先完成对小规模、缓存友好的输入数据子集的所有层计算。
• 识别并分组可优化层（特别是逐元素运算和池化层），这些层可在不改变最终输出的前提下对部分输入张量进行计算。
• 采用模块化运行时系统，支持可插拔的深度学习框架前端和硬件目标后端，实现透明集成。
• 应用堆叠技术，对连续的非卷积层在小数据块上进行处理，以最小化内存带宽压力并提升缓存复用率。
• 采用轻量级代码转换和内核融合策略，在保持数值正确性的同时维持高性能。
• 通过将相同的深度优先执行模型扩展至反向传播工作流，支持推理和未来训练的加速。

实验结果

研究问题

• RQ1 深度优先计算策略是否能显著改善深度神经网络推理中的数据局部性并减少缓存未命中？
• RQ2 基于数据分区的透明、非侵入式优化在标准 CPU 和 GPU 上能实现多大程度的性能提升？
• RQ3 哪些层类型适合深度优先处理，哪些层因数据依赖或计算模式而本质上不兼容？
• RQ4 BrainSlug 的性能提升在不同批量大小、网络架构和硬件平台上的表现如何变化？
• RQ5 该框架能否扩展以支持训练工作负载，并实现类似的性能提升？

主要发现

• BrainSlug 在多种深度神经网络模型上，CPU 硬件最高实现 41.1% 的加速，GPU 硬件最高实现 35.7% 的加速。
• 性能提升在批量大小为 8 及以上时最为显著，这在实践中很常见，例如 DenseNet 训练中常用的 32 批次大小。
• 该方法对用户透明，仅需极少代码修改——PyTorch 前端仅需 270 行 Python 代码和 438 行 C++ 代码，且不改变模型行为或架构。
• 卷积层和全连接层无法从深度优先方法中受益，前者因固有的数据重叠，后者因向量化操作中缓存复用性差。
• 该框架具有良好的可扩展性，其清晰的 API 设计支持轻松集成新的前端（如 TensorFlow）和后端（如 FPGA、向量处理器）。
• 未来工作表明，训练工作负载也可能实现类似的加速，目前正致力于将框架扩展以支持反向传播和更多硬件目标。