[论文解读] Auto-Vectorizing TensorFlow Graphs: Jacobians, Auto-Batching And Beyond

Ashish Agarwal, Igor Ganichev|arXiv (Cornell University)|Mar 8, 2019

Scientific Computing and Data Management参考文献 18被引用 6

一句话总结

本论文提出了一种针对TensorFlow高级数据流IR的静态循环向量化优化，通过一种新的符号并行for（pfor）抽象，实现了高效的自动批处理、逐样本梯度计算以及雅可比矩阵计算。该方法在GPU上相较于TensorFlow原生循环实现最高可提升60倍性能，相较于DyNet的动态批处理也超过38倍，尤其在大批次和高维输出场景下表现显著。

ABSTRACT

We propose a static loop vectorization optimization on top of high level dataflow IR used by frameworks like TensorFlow. A new statically vectorized parallel-for abstraction is provided on top of TensorFlow, and used for applications ranging from auto-batching and per-example gradients, to jacobian computation, optimized map functions and input pipeline optimization. We report huge speedups compared to both loop based implementations, as well as run-time batching adopted by the DyNet framework.

研究动机与目标

• 解决TensorFlow中基于循环的实现方式在自动批处理、逐样本梯度计算及雅可比矩阵计算等操作中的性能瓶颈问题。
• 通过实现静态、编译时向量化，减少DyNet等框架中动态批处理相关的运行时开销。
• 提升硬件利用率，支持机器学习模型中嵌套和动态控制流模式的高效执行。
• 提供一种通用、可复用的抽象，用于对TensorFlow高级IR中的任意循环进行向量化处理。
• 扩展TensorFlow原生能力，原生支持雅可比矩阵与海森矩阵的高效计算，而这些目前在性能优化上表现不佳或尚未被支持。

提出的方法

• 在图构建阶段，于TensorFlow高级数据流IR上引入一种新的静态向量化并行for（pfor）抽象。
• 通过将循环体融合为向量化操作，实现静态循环向量化，从而消除循环，适用于数据并行性极强的计算。
• 使用符号张量表示向量化循环迭代的输出，支持在嵌套控制流和动态形状中进行优化。
• 通过将前向传播封装在批量元素循环中，再对整个循环进行向量化，将pfor抽象应用于自动批处理模型。
• 通过在单一pfor结构内向量化前向与反向传播，实现逐样本梯度计算。
• 通过在每个输出元素上向量化梯度计算，将每个元素视为独立的标量计算，实现雅可比矩阵计算。

实验结果

研究问题

• RQ1 在TensorFlow高级IR中，静态循环向量化能否显著优于基于动态循环的实现方式，特别是在自动批处理与梯度计算方面？
• RQ2 pfor向量化在雅可比矩阵计算方面能带来多大性能提升，尤其是在非标量输出和复杂模型（如RNN）中？
• RQ3 pfor在不同批量大小与输出维度下，与DyNet的运行时批处理相比性能如何？
• RQ4 pfor当前实现中的性能瓶颈是什么，特别是在可变长度序列或嵌套循环场景下？
• RQ5 pfor能否通过向量化控制流实现此前不可行或低效的操作，如海森矩阵计算或输入流水线优化？

主要发现

• 在输出维度较大（128）时，pfor实现的雅可比矩阵计算相比TensorFlow原生tf.while_loop最高可提升60倍，尤其在VGG16和LSTM模型上表现显著。
• 在LSTM模型上计算逐样本梯度时，pfor在批量大小为256时相比tf.while_loop性能提升达38倍，展现出更优的GPU利用率。
• 在MNIST数据集上，pfor相比tf.while_loop展现出一致的吞吐量提升，尤其在大批次下优势明显，得益于更优的向量化与内存访问模式。
• 随着输出维度的增加，pfor与tf.while_loop之间的性能差距进一步扩大，表明该向量化方法具有良好的可扩展性。
• LSTM模型的相对加速比相对较低，原因在于序列长度可变以及tf.while_loop向量化带来的残余开销，提示在处理动态形状方面仍有优化空间。
• 该方法使此前不被支持的操作成为可能，如动态RNN的雅可比矩阵与非标量输出的海森矩阵计算，为新研究方向打开大门。